lockbot 2.0.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (54) hide show
  1. lockbot-2.0.0/LICENSE +21 -0
  2. lockbot-2.0.0/MANIFEST.in +3 -0
  3. lockbot-2.0.0/PKG-INFO +193 -0
  4. lockbot-2.0.0/README.md +153 -0
  5. lockbot-2.0.0/pyproject.toml +78 -0
  6. lockbot-2.0.0/python/lockbot/__init__.py +0 -0
  7. lockbot-2.0.0/python/lockbot/backend/__init__.py +0 -0
  8. lockbot-2.0.0/python/lockbot/backend/app/__init__.py +0 -0
  9. lockbot-2.0.0/python/lockbot/backend/app/admin/__init__.py +0 -0
  10. lockbot-2.0.0/python/lockbot/backend/app/admin/router.py +224 -0
  11. lockbot-2.0.0/python/lockbot/backend/app/auth/__init__.py +0 -0
  12. lockbot-2.0.0/python/lockbot/backend/app/auth/dependencies.py +72 -0
  13. lockbot-2.0.0/python/lockbot/backend/app/auth/models.py +24 -0
  14. lockbot-2.0.0/python/lockbot/backend/app/auth/router.py +140 -0
  15. lockbot-2.0.0/python/lockbot/backend/app/auth/schemas.py +70 -0
  16. lockbot-2.0.0/python/lockbot/backend/app/bots/__init__.py +0 -0
  17. lockbot-2.0.0/python/lockbot/backend/app/bots/encryption.py +57 -0
  18. lockbot-2.0.0/python/lockbot/backend/app/bots/manager.py +106 -0
  19. lockbot-2.0.0/python/lockbot/backend/app/bots/models.py +55 -0
  20. lockbot-2.0.0/python/lockbot/backend/app/bots/router.py +908 -0
  21. lockbot-2.0.0/python/lockbot/backend/app/bots/schemas.py +83 -0
  22. lockbot-2.0.0/python/lockbot/backend/app/bots/webhook_handler.py +70 -0
  23. lockbot-2.0.0/python/lockbot/backend/app/config.py +35 -0
  24. lockbot-2.0.0/python/lockbot/backend/app/database.py +24 -0
  25. lockbot-2.0.0/python/lockbot/backend/app/logs/__init__.py +0 -0
  26. lockbot-2.0.0/python/lockbot/backend/app/main.py +234 -0
  27. lockbot-2.0.0/python/lockbot/core/__init__.py +0 -0
  28. lockbot-2.0.0/python/lockbot/core/base_bot.py +193 -0
  29. lockbot-2.0.0/python/lockbot/core/bot_instance.py +49 -0
  30. lockbot-2.0.0/python/lockbot/core/config.py +325 -0
  31. lockbot-2.0.0/python/lockbot/core/device_bot.py +522 -0
  32. lockbot-2.0.0/python/lockbot/core/device_usage_alert.py +56 -0
  33. lockbot-2.0.0/python/lockbot/core/device_usage_utils.py +187 -0
  34. lockbot-2.0.0/python/lockbot/core/entry.py +81 -0
  35. lockbot-2.0.0/python/lockbot/core/env.py +11 -0
  36. lockbot-2.0.0/python/lockbot/core/handler.py +125 -0
  37. lockbot-2.0.0/python/lockbot/core/i18n/__init__.py +56 -0
  38. lockbot-2.0.0/python/lockbot/core/i18n/en.py +144 -0
  39. lockbot-2.0.0/python/lockbot/core/i18n/zh.py +137 -0
  40. lockbot-2.0.0/python/lockbot/core/io.py +217 -0
  41. lockbot-2.0.0/python/lockbot/core/message_adapter.py +49 -0
  42. lockbot-2.0.0/python/lockbot/core/msg_utils.py +114 -0
  43. lockbot-2.0.0/python/lockbot/core/node_bot.py +387 -0
  44. lockbot-2.0.0/python/lockbot/core/platforms/__init__.py +0 -0
  45. lockbot-2.0.0/python/lockbot/core/platforms/infoflow.py +119 -0
  46. lockbot-2.0.0/python/lockbot/core/queue_bot.py +470 -0
  47. lockbot-2.0.0/python/lockbot/core/request.py +110 -0
  48. lockbot-2.0.0/python/lockbot/core/utils.py +86 -0
  49. lockbot-2.0.0/python/lockbot.egg-info/PKG-INFO +193 -0
  50. lockbot-2.0.0/python/lockbot.egg-info/SOURCES.txt +52 -0
  51. lockbot-2.0.0/python/lockbot.egg-info/dependency_links.txt +1 -0
  52. lockbot-2.0.0/python/lockbot.egg-info/requires.txt +16 -0
  53. lockbot-2.0.0/python/lockbot.egg-info/top_level.txt +1 -0
  54. lockbot-2.0.0/setup.cfg +4 -0
lockbot-2.0.0/LICENSE ADDED
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2026 Jianbang Yang
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
@@ -0,0 +1,3 @@
1
+ include README.md
2
+ include LICENSE
3
+ recursive-include python/lockbot/core/i18n *.py
lockbot-2.0.0/PKG-INFO ADDED
@@ -0,0 +1,193 @@
1
+ Metadata-Version: 2.4
2
+ Name: lockbot
3
+ Version: 2.0.0
4
+ Summary: Cluster resource management bot for IM platforms
5
+ Author-email: Jianbang Yang <yangjianbang112@gmail.com>
6
+ License: MIT
7
+ Project-URL: Homepage, https://github.com/dynamicheart/lockbot
8
+ Project-URL: Repository, https://github.com/dynamicheart/lockbot
9
+ Project-URL: Issues, https://github.com/dynamicheart/lockbot/issues
10
+ Keywords: bot,cluster,resource-management,gpu,lock
11
+ Classifier: Development Status :: 4 - Beta
12
+ Classifier: Intended Audience :: Developers
13
+ Classifier: Intended Audience :: System Administrators
14
+ Classifier: License :: OSI Approved :: MIT License
15
+ Classifier: Programming Language :: Python :: 3
16
+ Classifier: Programming Language :: Python :: 3.10
17
+ Classifier: Programming Language :: Python :: 3.11
18
+ Classifier: Programming Language :: Python :: 3.12
19
+ Classifier: Topic :: System :: Clustering
20
+ Classifier: Topic :: Utilities
21
+ Requires-Python: >=3.10
22
+ Description-Content-Type: text/markdown
23
+ License-File: LICENSE
24
+ Requires-Dist: pycryptodome
25
+ Requires-Dist: requests
26
+ Requires-Dist: flask
27
+ Requires-Dist: six
28
+ Requires-Dist: fastapi
29
+ Requires-Dist: uvicorn[standard]
30
+ Requires-Dist: sqlalchemy
31
+ Requires-Dist: pyjwt
32
+ Requires-Dist: cryptography
33
+ Requires-Dist: bcrypt
34
+ Requires-Dist: python-multipart
35
+ Requires-Dist: httpx
36
+ Provides-Extra: dev
37
+ Requires-Dist: pytest>=7.0; extra == "dev"
38
+ Requires-Dist: ruff>=0.1.0; extra == "dev"
39
+ Dynamic: license-file
40
+
41
+ # lockbot
42
+
43
+ Cluster resource management bot for IM platforms (e.g., Baidu InfoFlow).
44
+
45
+ Lock and unlock GPU devices, cluster nodes, and queue slots via chat commands.
46
+ Supports both standalone Flask deployment and a full platform mode with FastAPI + Vue.js frontend.
47
+
48
+ [中文文档](README_CN.md) | [Live Demo](https://dynamicheart.github.io/lockbot/)
49
+
50
+ ## Features
51
+
52
+ - **Device Lock Bot** — Lock/unlock individual GPUs or devices on a cluster
53
+ - **Node Lock Bot** — Lock/unlock entire cluster nodes
54
+ - **Queue Bot** — Manage a queue for resource allocation with booking and preemption
55
+ - **Platform Mode** — Web UI (Vue 3 + Element Plus) for managing multiple bots, user authentication (JWT), role-based access control, and real-time logs
56
+ - **State Persistence** — Bot state survives restarts (JSON file)
57
+ - **Bilingual** — English and Chinese UI and bot responses
58
+
59
+ ## Quick Start — Platform Mode (Recommended)
60
+
61
+ Full management platform with Web UI, multi-bot orchestration, user authentication, and admin panel.
62
+
63
+ 1. Install:
64
+
65
+ ```bash
66
+ pip install lockbot
67
+ ```
68
+
69
+ 2. Set environment variables:
70
+
71
+ ```bash
72
+ export JWT_SECRET="your-jwt-secret"
73
+ export ENCRYPTION_KEY="your-fernet-key" # python -c "from cryptography.fernet import Fernet; print(Fernet.generate_key().decode())"
74
+ export DEV_MODE="true" # dev mode, auto-create admin user
75
+ ```
76
+
77
+ 3. Start:
78
+
79
+ ```bash
80
+ # Backend
81
+ uvicorn lockbot.backend.app.main:app --host 0.0.0.0 --port 8000 --reload
82
+
83
+ # Frontend (another terminal)
84
+ cd frontend && npm install && npm run dev
85
+ ```
86
+
87
+ 4. Open `http://localhost:8000` in your browser.
88
+
89
+ ### Docker
90
+
91
+ ```bash
92
+ docker build -f docker/Dockerfile -t lockbot .
93
+
94
+ # Platform mode (default)
95
+ docker run -d -p 8000:8000 \
96
+ -e JWT_SECRET=your-secret \
97
+ -e ENCRYPTION_KEY=your-fernet-key \
98
+ -v lockbot-data:/app/python/lockbot/data \
99
+ -v lockbot-state:/data/bots \
100
+ lockbot
101
+ ```
102
+
103
+ ## Bot Configuration
104
+
105
+ | Key | Description | Default |
106
+ |---|---|---|
107
+ | `BOT_TYPE` | `DEVICE`, `NODE`, or `QUEUE` | (required) |
108
+ | `BOT_NAME` | Bot instance name | `demo_bot` |
109
+ | `CLUSTER_CONFIGS` | Cluster layout (dict or list) | `{}` |
110
+ | `TOKEN` | Bot signature verification token | `""` |
111
+ | `AESKEY` | Message decryption AES key | `""` |
112
+ | `WEBHOOK_URL` | Message webhook URL | `""` |
113
+ | `PORT` | Server listen port | `8090` |
114
+ | `DEFAULT_DURATION` | Default lock duration (seconds) | `7200` (2h) |
115
+ | `MAX_LOCK_DURATION` | Max lock duration (seconds) | `-1` (unlimited) |
116
+ | `EARLY_NOTIFY` | Notify before lock expiry | `false` |
117
+
118
+ See `python/lockbot/core/config.py` for the full configuration reference.
119
+
120
+ ## Commands
121
+
122
+ | Command | Description |
123
+ |---------|-------------|
124
+ | `lock <node> [duration]` | Exclusive lock (e.g., `lock gpu0 3d`, `lock node1 30m`) |
125
+ | `slock <node> [duration]` | Shared lock (multiple users) |
126
+ | `unlock <node>` / `free <node>` | Release a specific node |
127
+ | `unlock` / `free` | Release all your nodes |
128
+ | `kickout <node>` | Force release (admin) |
129
+ | `book <node> [duration]` | Queue: book a node for later |
130
+ | `take <node>` | Queue: take the current lock |
131
+ | `<node>` | Query current usage |
132
+ | `help` | Show usage |
133
+
134
+ ## Development
135
+
136
+ ```bash
137
+ # Install dev dependencies
138
+ pip install -e ".[dev]"
139
+
140
+ # Run tests
141
+ pytest
142
+
143
+ # Lint + format check
144
+ ruff check python/ tests/
145
+ ruff format --check python/ tests/
146
+ ```
147
+
148
+ ## Standalone Mode
149
+
150
+ Single-bot deployment with a lightweight Flask webhook server.
151
+
152
+ **Device Lock Bot** (per-GPU locking):
153
+
154
+ ```python
155
+ from lockbot.core.bot_instance import BotInstance
156
+ from lockbot.core.entry import create_app
157
+
158
+ instance = BotInstance("DEVICE", {
159
+ "BOT_NAME": "my-gpu-bot",
160
+ "WEBHOOK_URL": "https://your-webhook-url",
161
+ "TOKEN": "your-bot-token",
162
+ "AESKEY": "your-aes-key",
163
+ "CLUSTER_CONFIGS": {
164
+ "node0": ["A800", "A800", "H100"],
165
+ "node1": ["A800", "H100"],
166
+ },
167
+ })
168
+
169
+ app = create_app(bot=instance.bot, bot_name="my-gpu-bot", port=8000)
170
+ app.run(host="0.0.0.0", port=8000)
171
+ ```
172
+
173
+ **Node Lock Bot / Queue Bot** (per-node locking or queue scheduling):
174
+
175
+ ```python
176
+ from lockbot.core.bot_instance import BotInstance
177
+ from lockbot.core.entry import create_app
178
+
179
+ instance = BotInstance("NODE", { # or "QUEUE" for queue scheduling
180
+ "BOT_NAME": "my-node-bot",
181
+ "WEBHOOK_URL": "https://your-webhook-url",
182
+ "TOKEN": "your-bot-token",
183
+ "AESKEY": "your-aes-key",
184
+ "CLUSTER_CONFIGS": ["node0", "node1", "node2", "node3"],
185
+ })
186
+
187
+ app = create_app(bot=instance.bot, bot_name="my-node-bot", port=8000)
188
+ app.run(host="0.0.0.0", port=8000)
189
+ ```
190
+
191
+ ## License
192
+
193
+ MIT
@@ -0,0 +1,153 @@
1
+ # lockbot
2
+
3
+ Cluster resource management bot for IM platforms (e.g., Baidu InfoFlow).
4
+
5
+ Lock and unlock GPU devices, cluster nodes, and queue slots via chat commands.
6
+ Supports both standalone Flask deployment and a full platform mode with FastAPI + Vue.js frontend.
7
+
8
+ [中文文档](README_CN.md) | [Live Demo](https://dynamicheart.github.io/lockbot/)
9
+
10
+ ## Features
11
+
12
+ - **Device Lock Bot** — Lock/unlock individual GPUs or devices on a cluster
13
+ - **Node Lock Bot** — Lock/unlock entire cluster nodes
14
+ - **Queue Bot** — Manage a queue for resource allocation with booking and preemption
15
+ - **Platform Mode** — Web UI (Vue 3 + Element Plus) for managing multiple bots, user authentication (JWT), role-based access control, and real-time logs
16
+ - **State Persistence** — Bot state survives restarts (JSON file)
17
+ - **Bilingual** — English and Chinese UI and bot responses
18
+
19
+ ## Quick Start — Platform Mode (Recommended)
20
+
21
+ Full management platform with Web UI, multi-bot orchestration, user authentication, and admin panel.
22
+
23
+ 1. Install:
24
+
25
+ ```bash
26
+ pip install lockbot
27
+ ```
28
+
29
+ 2. Set environment variables:
30
+
31
+ ```bash
32
+ export JWT_SECRET="your-jwt-secret"
33
+ export ENCRYPTION_KEY="your-fernet-key" # python -c "from cryptography.fernet import Fernet; print(Fernet.generate_key().decode())"
34
+ export DEV_MODE="true" # dev mode, auto-create admin user
35
+ ```
36
+
37
+ 3. Start:
38
+
39
+ ```bash
40
+ # Backend
41
+ uvicorn lockbot.backend.app.main:app --host 0.0.0.0 --port 8000 --reload
42
+
43
+ # Frontend (another terminal)
44
+ cd frontend && npm install && npm run dev
45
+ ```
46
+
47
+ 4. Open `http://localhost:8000` in your browser.
48
+
49
+ ### Docker
50
+
51
+ ```bash
52
+ docker build -f docker/Dockerfile -t lockbot .
53
+
54
+ # Platform mode (default)
55
+ docker run -d -p 8000:8000 \
56
+ -e JWT_SECRET=your-secret \
57
+ -e ENCRYPTION_KEY=your-fernet-key \
58
+ -v lockbot-data:/app/python/lockbot/data \
59
+ -v lockbot-state:/data/bots \
60
+ lockbot
61
+ ```
62
+
63
+ ## Bot Configuration
64
+
65
+ | Key | Description | Default |
66
+ |---|---|---|
67
+ | `BOT_TYPE` | `DEVICE`, `NODE`, or `QUEUE` | (required) |
68
+ | `BOT_NAME` | Bot instance name | `demo_bot` |
69
+ | `CLUSTER_CONFIGS` | Cluster layout (dict or list) | `{}` |
70
+ | `TOKEN` | Bot signature verification token | `""` |
71
+ | `AESKEY` | Message decryption AES key | `""` |
72
+ | `WEBHOOK_URL` | Message webhook URL | `""` |
73
+ | `PORT` | Server listen port | `8090` |
74
+ | `DEFAULT_DURATION` | Default lock duration (seconds) | `7200` (2h) |
75
+ | `MAX_LOCK_DURATION` | Max lock duration (seconds) | `-1` (unlimited) |
76
+ | `EARLY_NOTIFY` | Notify before lock expiry | `false` |
77
+
78
+ See `python/lockbot/core/config.py` for the full configuration reference.
79
+
80
+ ## Commands
81
+
82
+ | Command | Description |
83
+ |---------|-------------|
84
+ | `lock <node> [duration]` | Exclusive lock (e.g., `lock gpu0 3d`, `lock node1 30m`) |
85
+ | `slock <node> [duration]` | Shared lock (multiple users) |
86
+ | `unlock <node>` / `free <node>` | Release a specific node |
87
+ | `unlock` / `free` | Release all your nodes |
88
+ | `kickout <node>` | Force release (admin) |
89
+ | `book <node> [duration]` | Queue: book a node for later |
90
+ | `take <node>` | Queue: take the current lock |
91
+ | `<node>` | Query current usage |
92
+ | `help` | Show usage |
93
+
94
+ ## Development
95
+
96
+ ```bash
97
+ # Install dev dependencies
98
+ pip install -e ".[dev]"
99
+
100
+ # Run tests
101
+ pytest
102
+
103
+ # Lint + format check
104
+ ruff check python/ tests/
105
+ ruff format --check python/ tests/
106
+ ```
107
+
108
+ ## Standalone Mode
109
+
110
+ Single-bot deployment with a lightweight Flask webhook server.
111
+
112
+ **Device Lock Bot** (per-GPU locking):
113
+
114
+ ```python
115
+ from lockbot.core.bot_instance import BotInstance
116
+ from lockbot.core.entry import create_app
117
+
118
+ instance = BotInstance("DEVICE", {
119
+ "BOT_NAME": "my-gpu-bot",
120
+ "WEBHOOK_URL": "https://your-webhook-url",
121
+ "TOKEN": "your-bot-token",
122
+ "AESKEY": "your-aes-key",
123
+ "CLUSTER_CONFIGS": {
124
+ "node0": ["A800", "A800", "H100"],
125
+ "node1": ["A800", "H100"],
126
+ },
127
+ })
128
+
129
+ app = create_app(bot=instance.bot, bot_name="my-gpu-bot", port=8000)
130
+ app.run(host="0.0.0.0", port=8000)
131
+ ```
132
+
133
+ **Node Lock Bot / Queue Bot** (per-node locking or queue scheduling):
134
+
135
+ ```python
136
+ from lockbot.core.bot_instance import BotInstance
137
+ from lockbot.core.entry import create_app
138
+
139
+ instance = BotInstance("NODE", { # or "QUEUE" for queue scheduling
140
+ "BOT_NAME": "my-node-bot",
141
+ "WEBHOOK_URL": "https://your-webhook-url",
142
+ "TOKEN": "your-bot-token",
143
+ "AESKEY": "your-aes-key",
144
+ "CLUSTER_CONFIGS": ["node0", "node1", "node2", "node3"],
145
+ })
146
+
147
+ app = create_app(bot=instance.bot, bot_name="my-node-bot", port=8000)
148
+ app.run(host="0.0.0.0", port=8000)
149
+ ```
150
+
151
+ ## License
152
+
153
+ MIT
@@ -0,0 +1,78 @@
1
+ [build-system]
2
+ requires = ["setuptools>=68.0", "wheel"]
3
+ build-backend = "setuptools.build_meta"
4
+
5
+ [project]
6
+ name = "lockbot"
7
+ version = "2.0.0"
8
+ description = "Cluster resource management bot for IM platforms"
9
+ readme = "README.md"
10
+ license = {text = "MIT"}
11
+ requires-python = ">=3.10"
12
+ authors = [
13
+ {name = "Jianbang Yang", email = "yangjianbang112@gmail.com"},
14
+ ]
15
+ keywords = ["bot", "cluster", "resource-management", "gpu", "lock"]
16
+ classifiers = [
17
+ "Development Status :: 4 - Beta",
18
+ "Intended Audience :: Developers",
19
+ "Intended Audience :: System Administrators",
20
+ "License :: OSI Approved :: MIT License",
21
+ "Programming Language :: Python :: 3",
22
+ "Programming Language :: Python :: 3.10",
23
+ "Programming Language :: Python :: 3.11",
24
+ "Programming Language :: Python :: 3.12",
25
+ "Topic :: System :: Clustering",
26
+ "Topic :: Utilities",
27
+ ]
28
+
29
+ dependencies = [
30
+ # Core (standalone bot)
31
+ "pycryptodome",
32
+ "requests",
33
+ "flask",
34
+ "six",
35
+ # Backend (platform mode)
36
+ "fastapi",
37
+ "uvicorn[standard]",
38
+ "sqlalchemy",
39
+ "pyjwt",
40
+ "cryptography",
41
+ "bcrypt",
42
+ "python-multipart",
43
+ "httpx",
44
+ ]
45
+
46
+ [project.optional-dependencies]
47
+ dev = [
48
+ "pytest>=7.0",
49
+ "ruff>=0.1.0",
50
+ ]
51
+
52
+ [project.urls]
53
+ Homepage = "https://github.com/dynamicheart/lockbot"
54
+ Repository = "https://github.com/dynamicheart/lockbot"
55
+ Issues = "https://github.com/dynamicheart/lockbot/issues"
56
+
57
+ [tool.setuptools.packages.find]
58
+ where = ["python"]
59
+
60
+ [tool.setuptools.package-data]
61
+ lockbot = ["core/i18n/*.py"]
62
+
63
+ [tool.ruff]
64
+ target-version = "py310"
65
+ line-length = 120
66
+
67
+ [tool.ruff.lint]
68
+ select = ["E", "F", "W", "I", "UP", "B", "SIM"]
69
+ ignore = ["SIM108"]
70
+
71
+ [tool.ruff.lint.per-file-ignores]
72
+ "python/lockbot/backend/**" = ["B008"] # FastAPI Depends() in defaults is idiomatic
73
+
74
+ [tool.ruff.format]
75
+ quote-style = "double"
76
+
77
+ [tool.pytest.ini_options]
78
+ testpaths = ["tests"]
File without changes
File without changes
File without changes
@@ -0,0 +1,224 @@
1
+ """
2
+ Admin API routes.
3
+ """
4
+
5
+ import contextlib
6
+ import os
7
+ import shutil
8
+ import tempfile
9
+ from datetime import datetime
10
+
11
+ from fastapi import APIRouter, Depends, HTTPException
12
+ from fastapi.responses import FileResponse
13
+ from sqlalchemy.orm import Session
14
+ from starlette.background import BackgroundTask
15
+
16
+ from lockbot.backend.app.auth.dependencies import can_manage_user, require_admin, require_super_admin
17
+ from lockbot.backend.app.auth.models import User
18
+ from lockbot.backend.app.auth.router import _generate_password, _hash_password
19
+ from lockbot.backend.app.auth.schemas import (
20
+ AdminCreateUser,
21
+ AdminEditUser,
22
+ PasswordResetOut,
23
+ UserOut,
24
+ )
25
+ from lockbot.backend.app.bots.models import Bot
26
+ from lockbot.backend.app.database import get_db
27
+
28
+ router = APIRouter(prefix="/api/admin", tags=["admin"])
29
+
30
+ ADMIN_VISIBLE_ROLES = ("admin", "user")
31
+ SUPER_ADMIN_VISIBLE_ROLES = ("super_admin", "admin", "user")
32
+ ADMIN_CREATABLE_ROLES = ("user",)
33
+
34
+
35
+ @router.post("/users", response_model=PasswordResetOut, status_code=201)
36
+ def admin_create_user(
37
+ body: AdminCreateUser,
38
+ operator: User = Depends(require_admin),
39
+ db: Session = Depends(get_db),
40
+ ):
41
+ """Admin creates a user. Super_admin can create admin/user; admin can only create user."""
42
+ # Non-super_admin can only create regular users
43
+ if operator.role != "super_admin" and body.role not in ADMIN_CREATABLE_ROLES:
44
+ raise HTTPException(status_code=403, detail="Cannot create user with this role")
45
+
46
+ if body.role not in SUPER_ADMIN_VISIBLE_ROLES:
47
+ raise HTTPException(status_code=400, detail=f"Invalid role, must be one of {SUPER_ADMIN_VISIBLE_ROLES}")
48
+
49
+ exists = db.query(User).filter(User.username == body.username).first()
50
+ if exists:
51
+ raise HTTPException(status_code=409, detail="Username already taken")
52
+ dup_email = db.query(User).filter(User.email == body.email).first()
53
+ if dup_email:
54
+ raise HTTPException(status_code=409, detail="Email already taken")
55
+
56
+ raw_password = _generate_password()
57
+ user = User(
58
+ username=body.username,
59
+ email=body.email,
60
+ password_hash=_hash_password(raw_password),
61
+ role=body.role,
62
+ max_running_bots=body.max_running_bots,
63
+ must_change_password=True,
64
+ )
65
+ db.add(user)
66
+ db.commit()
67
+ db.refresh(user)
68
+ return PasswordResetOut(id=user.id, username=user.username, new_password=raw_password)
69
+
70
+
71
+ @router.get("/users", response_model=list[UserOut])
72
+ def list_users(
73
+ operator: User = Depends(require_admin),
74
+ db: Session = Depends(get_db),
75
+ ):
76
+ """List users visible to the operator.
77
+ Super_admin sees all; admin sees only users (not admins/super_admins)."""
78
+ if operator.role == "super_admin":
79
+ return db.query(User).all()
80
+
81
+ # admin: see regular users + self
82
+ return db.query(User).filter((User.role == "user") | (User.id == operator.id)).all()
83
+
84
+
85
+ @router.put("/users/{user_id}", response_model=UserOut)
86
+ def admin_edit_user(
87
+ user_id: int,
88
+ body: AdminEditUser,
89
+ operator: User = Depends(require_admin),
90
+ db: Session = Depends(get_db),
91
+ ):
92
+ """Admin edits user profile."""
93
+ target = db.get(User, user_id)
94
+ if not target:
95
+ raise HTTPException(status_code=404, detail="User not found")
96
+
97
+ # Cannot edit users at or above your own level
98
+ if not can_manage_user(operator, target.role):
99
+ raise HTTPException(status_code=403, detail="Cannot manage this user")
100
+
101
+ if body.username is not None and body.username != target.username:
102
+ dup = db.query(User).filter(User.username == body.username).first()
103
+ if dup:
104
+ raise HTTPException(status_code=409, detail="Username already taken")
105
+ target.username = body.username
106
+
107
+ if body.email is not None and body.email != target.email:
108
+ dup = db.query(User).filter(User.email == body.email).first()
109
+ if dup:
110
+ raise HTTPException(status_code=409, detail="Email already taken")
111
+ target.email = body.email
112
+
113
+ if body.role is not None:
114
+ if body.role not in SUPER_ADMIN_VISIBLE_ROLES:
115
+ raise HTTPException(status_code=400, detail="Invalid role")
116
+ # Only super_admin can promote to admin
117
+ if body.role in ("admin", "super_admin") and operator.role != "super_admin":
118
+ raise HTTPException(status_code=403, detail="Only super admin can set this role")
119
+ target.role = body.role
120
+
121
+ if body.max_running_bots is not None:
122
+ target.max_running_bots = body.max_running_bots
123
+
124
+ db.commit()
125
+ db.refresh(target)
126
+ return target
127
+
128
+
129
+ @router.put("/users/{user_id}/max-bots")
130
+ def set_max_bots(
131
+ user_id: int,
132
+ body: dict,
133
+ operator: User = Depends(require_admin),
134
+ db: Session = Depends(get_db),
135
+ ):
136
+ user = db.get(User, user_id)
137
+ if not user:
138
+ raise HTTPException(status_code=404, detail="User not found")
139
+ if not can_manage_user(operator, user.role):
140
+ raise HTTPException(status_code=403, detail="Cannot manage this user")
141
+ user.max_running_bots = body["max_running_bots"]
142
+ db.commit()
143
+ db.refresh(user)
144
+ return {"id": user.id, "max_running_bots": user.max_running_bots}
145
+
146
+
147
+ @router.post("/users/{user_id}/reset-password", response_model=PasswordResetOut)
148
+ def admin_reset_password(
149
+ user_id: int,
150
+ operator: User = Depends(require_admin),
151
+ db: Session = Depends(get_db),
152
+ ):
153
+ """Admin resets a user's password. Returns the new plaintext password."""
154
+ target = db.get(User, user_id)
155
+ if not target:
156
+ raise HTTPException(status_code=404, detail="User not found")
157
+ if not can_manage_user(operator, target.role):
158
+ raise HTTPException(status_code=403, detail="Cannot manage this user")
159
+ raw_password = _generate_password()
160
+ target.password_hash = _hash_password(raw_password)
161
+ target.must_change_password = True
162
+ db.commit()
163
+ db.refresh(target)
164
+ return PasswordResetOut(id=target.id, username=target.username, new_password=raw_password)
165
+
166
+
167
+ @router.get("/bots")
168
+ def list_all_bots(
169
+ _admin: User = Depends(require_admin),
170
+ db: Session = Depends(get_db),
171
+ ):
172
+ rows = db.query(Bot, User.username).join(User, Bot.user_id == User.id).all()
173
+ return [
174
+ {c.name: getattr(bot, c.name) for c in bot.__table__.columns} | {"owner": username} for bot, username in rows
175
+ ]
176
+
177
+
178
+ @router.get("/stats")
179
+ def platform_stats(
180
+ _admin: User = Depends(require_admin),
181
+ db: Session = Depends(get_db),
182
+ ):
183
+ total_users = db.query(User).count()
184
+ bots = db.query(Bot).all()
185
+ return {
186
+ "totalUsers": total_users,
187
+ "totalBots": len(bots),
188
+ "running": sum(1 for b in bots if b.status == "running"),
189
+ "errors": sum(1 for b in bots if b.status == "error"),
190
+ }
191
+
192
+
193
+ @router.get("/backup")
194
+ def download_backup(
195
+ _admin: User = Depends(require_super_admin),
196
+ ):
197
+ """Download full SQLite database backup (super_admin only)."""
198
+ from lockbot.backend.app.config import DATABASE_URL
199
+
200
+ db_url = DATABASE_URL
201
+ if not db_url.startswith("sqlite:///"):
202
+ raise HTTPException(status_code=400, detail="Backup only supported for SQLite")
203
+
204
+ db_path = db_url[len("sqlite:///") :]
205
+ if not os.path.exists(db_path):
206
+ raise HTTPException(status_code=404, detail="Database file not found")
207
+
208
+ # Copy to temp file to avoid sending a mid-write database
209
+ fd, tmp_path = tempfile.mkstemp(suffix=".db")
210
+ os.close(fd)
211
+ try:
212
+ shutil.copy2(db_path, tmp_path)
213
+ except BaseException:
214
+ with contextlib.suppress(OSError):
215
+ os.unlink(tmp_path)
216
+ raise
217
+
218
+ filename = f"lockbot_backup_{datetime.now().strftime('%Y%m%d_%H%M%S')}.db"
219
+ return FileResponse(
220
+ path=tmp_path,
221
+ filename=filename,
222
+ media_type="application/x-sqlite3",
223
+ background=BackgroundTask(os.unlink, tmp_path),
224
+ )