diff --git a/core/schemas/user.py b/core/schemas/user.py index 0f853a0e3..305b96bc3 100644 --- a/core/schemas/user.py +++ b/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 diff --git a/core/web/apiv2/auth.py b/core/web/apiv2/auth.py index 885ad350f..979f18940 100644 --- a/core/web/apiv2/auth.py +++ b/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)): diff --git a/extras/docker/scripts/docker-entrypoint.sh b/extras/docker/scripts/docker-entrypoint.sh index 5325815c3..51114346a 100755 --- a/extras/docker/scripts/docker-entrypoint.sh +++ b/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 "$@" diff --git a/extras/v1migrate/README.md b/extras/v1migrate/README.md new file mode 100644 index 000000000..410e39969 --- /dev/null +++ b/extras/v1migrate/README.md @@ -0,0 +1,8 @@ +# migrate from old yeti + + +## Deps + +``` +pip install pymongo +``` diff --git a/extras/yetictl/__init__.py b/extras/yetictl/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/extras/yetictl/cli.py b/extras/yetictl/cli.py new file mode 100644 index 000000000..4d4beb80e --- /dev/null +++ b/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() diff --git a/poetry.lock b/poetry.lock index 0c28d4dbc..7f86bc330 100644 --- a/poetry.lock +++ b/poetry.lock @@ -2263,4 +2263,4 @@ testing = ["big-O", "jaraco.functools", "jaraco.itertools", "more-itertools", "p [metadata] lock-version = "2.0" python-versions = ">=3.10,<3.12" -content-hash = "5e9d430887e4febe4322978966b00d13de1b6f880dd4ce3bbc5144127ef8de06" +content-hash = "9faeb36afc11f69a7c374efc6091fc07b3360bc3e3b763523ddf725717155535" diff --git a/pyproject.toml b/pyproject.toml index 1fe99ba1d..0697dc93c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -18,6 +18,7 @@ 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" @@ -25,6 +26,9 @@ 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"