User management script

core/schemas/user.py

@@ -16,6 +16,7 @@ def generate_api_key():
 class User(BaseModel, database_arango.ArangoYetiConnector):
 
     _collection_name: ClassVar[str] = "users"
+    _type_filter: ClassVar[None] = None
 
     id: str | None = None
     username: str

core/web/apiv2/auth.py

@@ -89,7 +89,9 @@ async def login(response: Response, form_data: OAuth2PasswordRequestForm = Depen
     if not YETI_AUTH:
         user = UserSensitive.find(username="yeti")
         if not user:
-            user = UserSensitive(username="yeti", admin=True).save()
+            user = UserSensitive(username="yeti", admin=True)
+            user.set_password("yeti")
+            user.save()
     else:
         user = UserSensitive.find(username=form_data.username)
         if not (user and user.verify_password(form_data.password)):

extras/docker/scripts/docker-entrypoint.sh

@@ -5,8 +5,11 @@ if [ "$1" = 'webserver' ]; then
     poetry run uvicorn core.web.webapp:app --reload --host 0.0.0.0
 elif  [ "$1" = 'tasks' ]; then
     poetry run celery -A core.taskmanager worker --loglevel=INFO --purge -B -P threads
+elif [ "$1" = 'create-user']; then
+    poetry run yetictl create-user "${@:2}"
+elif [ "$1" = 'reset-password']; then
+    poetry run yetictl reset-password "${@:2}"
 elif  [ "$1" = 'envshell' ]; then
     poetry shell
-fi
-
-exec "$@"
+else
+    exec "$@"

extras/v1migrate/README.md

@@ -0,0 +1,8 @@
+# migrate from old yeti
+
+
+## Deps
+
+```
+pip install pymongo
+```

extras/yetictl/cli.py

@@ -0,0 +1,65 @@
+import click
+
+from core.schemas.user import UserSensitive, User
+
+
+@click.group()
+def cli():
+    pass
+
+
+@cli.command()
+def list_users():
+    users = User.list()
+    for user in users:
+        click.echo(
+            f"Username: {user.username} | API key: {user.api_key} | Admin: {user.admin}"
+        )
+
+
+@cli.command()
+@click.argument("username")
+@click.argument("password")
+@click.option("--admin", is_flag=True, default=False)
+def create_user(username: str, password: str, admin: bool = False) -> None:
+    """Creates a new user in the system."""
+    user = UserSensitive.find(username=username)
+    if user:
+        raise RuntimeError(f"User with username {username} already exists")
+    user = UserSensitive(username=username, admin=admin)
+    user.set_password(password)
+    user.save()
+    click.echo(
+        f"User {username} succesfully created! API key: {username}:{user.api_key}"
+    )
+
+
+@cli.command()
+@click.argument("username")
+def delete_user(username: str) -> None:
+    """Deletes a user from the system."""
+    user = UserSensitive.find(username=username)
+    if not user:
+        raise RuntimeError(f"User with username {username} does not exist")
+    user.delete()
+    click.echo(f"User {username} succesfully deleted")
+
+
+@cli.command()
+@click.argument("username")
+@click.argument("new_password")
+def reset_password(username: str, new_password: str) -> None:
+    """Resets a user's password."""
+    user = UserSensitive.find(username=username)
+    if not user:
+        raise RuntimeError(f"User with username {username} could not be found")
+    user.set_password(new_password)
+    user.reset_api_key()
+    user.save()
+    click.echo(
+        f"Password for {username} succesfully reset. New API key: {user.api_key}"
+    )
+
+
+if __name__ == "__main__":
+    cli()

pyproject.toml

@@ -18,13 +18,17 @@ passlib = {extras = ["bcrypt"], version = "^1.7.4"}
 python-multipart = "^0.0.6"
 pandas = "^2.1.1"
 redis = "^5.0.0"
+click = "^8.1.7"
 
 [tool.poetry.group.dev.dependencies]
 pylint = "^2.16.1"
 black = "^22.12.0"
 mypy = "^1.0.0"
 httpx = "^0.23.3"
 
+[tool.poetry.scripts]
+yetictl = 'extras.yetictl.cli:cli'
+
 [tool.poetry.group.plugins.dependencies]
 pymisp = "^2.4.176"
 otxv2 = "^1.5.12"