NetworkStatus-Rabbit/manage.py at master · wingsrabbit/NetworkStatus-Rabbit · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
#!/usr/bin/env python3
"""Management CLI commands - create-admin, remove-admin, reset-password."""
import argparse
import sys

import bcrypt


def get_app():
    from server.app import create_app
    return create_app()


def create_admin(args):
    app = get_app()
    with app.app_context():
        from server.extensions import db
        from server.models.user import User

        if User.query.filter_by(username=args.username).first():
            print(f"Error: username '{args.username}' already exists")
            sys.exit(1)

        password_hash = bcrypt.hashpw(
            args.password.encode('utf-8'),
            bcrypt.gensalt()
        ).decode('utf-8')

        user = User(
            username=args.username,
            password_hash=password_hash,
            role='admin',
            created_by='system',
        )
        db.session.add(user)
        db.session.commit()
        print(f"Admin '{args.username}' created successfully")


def remove_admin(args):
    app = get_app()
    with app.app_context():
        from server.extensions import db
        from server.models.user import User

        user = User.query.filter_by(username=args.username).first()
        if not user:
            print(f"Error: user '{args.username}' not found")
            sys.exit(1)

        # Rule 9: ensure at least 1 admin remains
        admin_count = User.query.filter_by(role='admin').count()
        if user.role == 'admin' and admin_count <= 1:
            print("Error: cannot remove the last admin. System must have at least 1 admin.")
            sys.exit(1)

        db.session.delete(user)
        db.session.commit()
        print(f"User '{args.username}' removed successfully")


def reset_password(args):
    app = get_app()
    with app.app_context():
        from server.extensions import db
        from server.models.user import User

        user = User.query.filter_by(username=args.username).first()
        if not user:
            print(f"Error: user '{args.username}' not found")
            sys.exit(1)

        if not args.password:
            print("Error: --password is required")
            sys.exit(1)

        user.password_hash = bcrypt.hashpw(
            args.password.encode('utf-8'),
            bcrypt.gensalt()
        ).decode('utf-8')
        user.failed_login_count = 0
        user.locked_until = None
        db.session.commit()
        print(f"Password for '{args.username}' reset successfully")


def main():
    parser = argparse.ArgumentParser(description='NetworkStatus-Rabbit Management CLI')
    subparsers = parser.add_subparsers(dest='command', required=True)

    # create-admin
    p_create = subparsers.add_parser('create-admin', help='Create a new admin user')
    p_create.add_argument('--username', required=True, help='Admin username')
    p_create.add_argument('--password', required=True, help='Admin password')

    # remove-admin
    p_remove = subparsers.add_parser('remove-admin', help='Remove a user')
    p_remove.add_argument('--username', required=True, help='Username to remove')

    # reset-password
    p_reset = subparsers.add_parser('reset-password', help='Reset a user password')
    p_reset.add_argument('--username', required=True, help='Username')
    p_reset.add_argument('--password', required=True, help='New password')

    args = parser.parse_args()

    if args.command == 'create-admin':
        create_admin(args)
    elif args.command == 'remove-admin':
        remove_admin(args)
    elif args.command == 'reset-password':
        reset_password(args)


if __name__ == '__main__':
    main()