Check in current backend

packages/sshecret-backend/src/sshecret_backend/cli.py

@@ -6,12 +6,12 @@ from pathlib import Path
 from typing import cast
 from dotenv import load_dotenv
 import click
-from sqlmodel import Session, create_engine, select
+from sqlmodel import Session, col, func, select
 import uvicorn
 
-from .db import create_api_token
+from .db import get_engine, create_api_token
 
-from .models import Client, ClientSecret, ClientAccessPolicy, AuditLog, APIClient
+from .models import Client, ClientSecret, ClientAccessPolicy, AuditLog, APIClient, init_db
 from .settings import BackendSettings
 
 DEFAULT_LISTEN = "127.0.0.1"
@@ -21,6 +21,23 @@ WORKDIR = Path(os.getcwd())
 
 load_dotenv()
 
+def generate_token(settings: BackendSettings) -> str:
+    """Generate a token."""
+    engine = get_engine(settings.db_url)
+    init_db(engine)
+    with Session(engine) as session:
+        token = create_api_token(session, True)
+    return token
+
+def count_tokens(settings: BackendSettings) -> int:
+    """Count the amount of tokens created."""
+    engine = get_engine(settings.db_url)
+    init_db(engine)
+    with Session(engine) as session:
+        count = session.exec(select(func.count("*")).select_from(APIClient)).one()
+
+    return count
+
 
 @click.group()
 @click.option("--database", help="Path to the sqlite database file.")
@@ -28,11 +45,19 @@ load_dotenv()
 def cli(ctx: click.Context, database: str) -> None:
     """CLI group."""
     if database:
-        # Hopefully it's enough to set the environment variable as so.
-        settings = BackendSettings(db_url=f"sqlite:///{Path(database).absolute()}")
+        settings = BackendSettings(database=str(Path(database).absolute()))
     else:
         settings = BackendSettings()
 
+
+    if settings.generate_initial_tokens:
+        if count_tokens(settings) == 0:
+            click.echo("Creating initial tokens for admin and sshd.")
+            admin_token = generate_token(settings)
+            sshd_token = generate_token(settings)
+            click.echo(f"Admin token: {admin_token}")
+            click.echo(f"SSHD token: {sshd_token}")
+
     ctx.obj = settings
 
 
@@ -41,9 +66,7 @@ def cli(ctx: click.Context, database: str) -> None:
 def cli_generate_token(ctx: click.Context) -> None:
     """Generate a token."""
     settings = cast(BackendSettings, ctx.obj)
-    engine = create_engine(settings.db_url)
-    with Session(engine) as session:
-        token = create_api_token(session, True)
+    token = generate_token(settings)
     click.echo("Generated api token:")
     click.echo(token)
 
@@ -61,7 +84,7 @@ def cli_run(host: str, port: int, dev: bool, workers: int | None) -> None:
 def cli_repl(ctx: click.Context) -> None:
     """Run an interactive console."""
     settings = cast(BackendSettings, ctx.obj)
-    engine = create_engine(settings.db_url)
+    engine = get_engine(settings.db_url, True)
 
     with Session(engine) as session:
         locals = {