feat!: Add Tracecat workspaces (TracecatHQ#302)

Co-authored-by: Daryl Lim <[email protected]>

.gitignore

@@ -62,8 +62,6 @@ dist/
 downloads/
 eggs/
 .eggs/
-lib/
-lib64/
 parts/
 sdist/
 var/
@@ -74,6 +72,14 @@ share/python-wheels/
 *.egg
 MANIFEST
 
+# Ignore Python-related lib directories
+/lib/
+/*/lib/
+/lib64/
+
+# But don't ignore the frontend src/lib directory
+!/frontend/src/lib/
+
 # PyInstaller
 #  Usually these files are written by a python script from a template
 #  before PyInstaller builds the exe, so as to inject date/other infos into it.
@@ -206,6 +212,3 @@ cython_debug/
 #  and can be added to the global gitignore or merged into this file.  For a more nuclear
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 #.idea/
-
-# Frontend
-!frontend/src/lib/*

cli/tracecat_cli/auth.py

@@ -4,9 +4,9 @@
 import rich
 import typer
 
-from ._client import Client
-from ._config import config
-from ._utils import delete_cookies, write_cookies
+from .client import Client
+from .config import config, manager
+from .utils import pprint_json
 
 app = typer.Typer(no_args_is_help=True, help="Authentication")
 
@@ -33,20 +33,26 @@ def login(
         )
         response.raise_for_status()
     # Convert cookies to a dictionary
-    write_cookies(response.cookies, config.cookies_path)
+    manager.write_cookies(response.cookies)
 
     rich.print(
-        f"[green]Login successful. Cookies saved to {config.cookies_path}[/green]"
+        f"[green]Login successful. Cookies saved to {config.config_path}[/green]"
     )
 
 
 @app.command(help="Get the current user")
-def whoami():
+def whoami(as_json: bool = typer.Option(False, "--json", help="Output as JSON")):
     """Get the current user."""
     with Client() as client:
         response = client.get("/users/me")
         response.raise_for_status()
-        rich.print(response.json())
+        user = response.json()
+    if as_json:
+        pprint_json(user)
+    else:
+        rich.print(f"Username: {user['email']}")
+        rich.print(f"First Name: {user['first_name']}")
+        rich.print(f"Last Name: {user['last_name']}")
 
 
 @app.command(
@@ -58,5 +64,5 @@ def logout():
         response = client.post("/auth/logout")
         response.raise_for_status()
     # Convert cookies to a dictionary
-    delete_cookies(config.cookies_path)
+    manager.delete_cookies()
     rich.print("[green]Logout successful[/green]")

cli/tracecat_cli/_client.py → cli/tracecat_cli/client.py

@@ -1,23 +1,32 @@
 import json
+import logging
 from typing import Any
 
 import httpx
 import rich
 import typer
 
-from ._config import config
-from ._utils import read_cookies
+from .config import config, manager
+
+logger = logging.getLogger(__name__)
 
 
 def _aclient():
+    workspace = manager.get_workspace()
+    logger.info(f"workspace: {workspace}")
     return httpx.AsyncClient(
-        base_url=config.api_url, cookies=read_cookies(config.cookies_path)
+        base_url=config.api_url,
+        cookies=manager.read_cookies(),
+        params={"workspace_id": workspace["id"]} if workspace else None,
     )
 
 
 def _client():
+    workspace = manager.get_workspace()
     return httpx.Client(
-        base_url=config.api_url, cookies=read_cookies(config.cookies_path)
+        base_url=config.api_url,
+        cookies=manager.read_cookies(),
+        params={"workspace_id": workspace["id"]} if workspace else None,
     )
 
 

cli/tracecat_cli/config.py

@@ -0,0 +1,106 @@
+import json
+import os
+import uuid
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import TypedDict
+
+import httpx
+from dotenv import find_dotenv, load_dotenv
+
+load_dotenv(find_dotenv())
+
+
+@dataclass
+class Role:
+    type: str
+    user_id: uuid.UUID | None
+    service_id: str
+
+
+class Workspace(TypedDict):
+    id: str
+    name: str
+
+
+# In reality we should use the user's id from config.toml
+@dataclass(frozen=True)
+class Config:
+    role: Role = field(
+        default_factory=lambda: Role(
+            type="service", user_id=uuid.UUID(int=0), service_id="tracecat-cli"
+        )
+    )
+    jwt_token: str = field(default="super-secret-jwt-token")
+    docs_path: Path = field(default_factory=lambda: Path("docs"))
+    docs_api_group: str = field(default="API Documentation")
+    docs_api_pages_group: str = field(default="Reference")
+    api_url: str = field(
+        default=os.getenv("TRACECAT__PUBLIC_API_URL", "http://localhost/api")
+    )
+    config_path: Path = field(default=Path.home() / ".tracecat_cli_config.json")
+
+
+config = Config()
+
+
+class ConfigFileManager:
+    def __init__(self, path: Path) -> None:
+        self.path = path
+
+    def __repr__(self) -> str:
+        data = self._read_config()
+        return f"ConfigFileManager(path={self.path}, data={data})"
+
+    def write_cookies(self, cookies: httpx.Cookies) -> None:
+        """Write cookies to config."""
+        cfg = self._read_config()
+        cfg["cookies"] = dict(cookies)
+        self._write_config(cfg)
+
+    def read_cookies(self) -> httpx.Cookies:
+        """Read cookies from config."""
+        cfg = self._read_config()
+        return httpx.Cookies(cfg.get("cookies", {}))
+
+    def delete_cookies(self) -> None:
+        """Delete cookies from config."""
+        cfg = self._read_config()
+        cfg.pop("cookies", None)
+        self._write_config(cfg)
+
+    def set_workspace(self, workspace_id: uuid.UUID, workspace_name: str) -> Workspace:
+        """Set the workspace ID in the configuration."""
+        cfg = self._read_config()
+        workspace = Workspace(id=str(workspace_id), name=workspace_name)
+        cfg["workspace"] = workspace
+        self._write_config(cfg)
+        return workspace
+
+    def get_workspace(self) -> Workspace | None:
+        """Get the workspace ID from the configuration."""
+        cfg = self._read_config()
+        workspace = cfg.get("workspace")
+        return Workspace(**workspace) if workspace else None
+
+    def reset_workspace(self) -> None:
+        """Remove the workspace ID from the configuration."""
+        cfg = self._read_config()
+        cfg.pop("workspace", None)
+        self._write_config(cfg)
+
+    def _read_config(self) -> dict[str, str]:
+        """Read configuration from file."""
+        try:
+            with self.path.open() as f:
+                return json.load(f)
+        except (FileNotFoundError, json.JSONDecodeError):
+            return {}
+
+    def _write_config(self, config: dict[str, str]) -> None:
+        """Write configuration to file."""
+        with self.path.open(mode="w") as f:
+            json.dump(config, f, indent=2)
+
+
+manager = ConfigFileManager(path=config.config_path)

cli/tracecat_cli/dev.py

@@ -16,9 +16,9 @@
 import yaml
 from pydantic import BaseModel
 
-from ._client import Client
-from ._config import config
-from ._utils import read_input
+from .client import Client
+from .config import config
+from .utils import read_input
 
 app = typer.Typer(no_args_is_help=True, help="Dev tools.")
 

cli/tracecat_cli/main.py

@@ -1,7 +1,7 @@
 import typer
 from dotenv import find_dotenv, load_dotenv
 
-from . import auth, dev, schedule, secret, workflow
+from . import auth, dev, schedule, secret, workflow, workspace
 
 load_dotenv(find_dotenv())
 app = typer.Typer(no_args_is_help=True, pretty_exceptions_show_locals=False)
@@ -28,6 +28,7 @@ def tracecat(
 app.add_typer(secret.app, name="secret")
 app.add_typer(schedule.app, name="schedule")
 app.add_typer(auth.app, name="auth")
+app.add_typer(workspace.app, name="workspace")
 
 if __name__ == "__main__":
     typer.run(app)

cli/tracecat_cli/schedule.py

@@ -7,8 +7,8 @@
 import typer
 from rich.console import Console
 
-from ._client import Client
-from ._utils import dynamic_table, read_input
+from .client import Client
+from .utils import dynamic_table, read_input
 
 app = typer.Typer(no_args_is_help=True, help="Manage schedules.")
 

cli/tracecat_cli/secret.py

@@ -1,13 +1,12 @@
-import asyncio
 from typing import TypedDict
 
 import orjson
 import rich
 import typer
 from rich.console import Console
 
-from ._client import Client
-from ._utils import dynamic_table
+from .client import Client
+from .utils import dynamic_table, pprint_json
 
 app = typer.Typer(no_args_is_help=True, help="Manage secrets.")
 
@@ -61,8 +60,7 @@ def list_secrets(
         result = Client.handle_response(res)
 
     if as_json:
-        out = orjson.dumps(result, option=orjson.OPT_INDENT_2).decode()
-        rich.print(out)
+        pprint_json(result)
     elif not result:
         rich.print("[cyan]No secrets found[/cyan]")
     else:
@@ -86,13 +84,17 @@ def delete(
         rich.print("Aborted")
         return
 
-    async def _delete():
-        async with Client() as client, asyncio.TaskGroup() as tg:
+    try:
+        with Client() as client:
             for name in secret_names:
-                tg.create_task(client.delete(f"/secrets/{name}"))
-
-    asyncio.run(_delete())
-    rich.print("[green]Secret deleted successfully![/green]")
+                get_response = client.get(f"/secrets/{name}")
+                secret = Client.handle_response(get_response)
+                del_response = client.delete(f"/secrets/{secret['id']}")
+                del_response.raise_for_status()
+        rich.print("[green]Secret deleted successfully![/green]")
+    except Exception as e:
+        rich.print(f"[red]Error: {e}[/red]")
+        return
 
 
 @app.command(no_args_is_help=True, help="Update a secret.")

cli/tracecat_cli/_utils.py → cli/tracecat_cli/utils.py

@@ -1,8 +1,8 @@
-import json
 from pathlib import Path
+from typing import Any
 
-import httpx
 import orjson
+import rich
 import typer
 from rich.table import Table
 
@@ -40,25 +40,6 @@ def read_input(data: str) -> dict[str, str]:
         raise typer.BadParameter(f"Invalid JSON: {e}") from e
 
 
-def write_cookies(cookies: httpx.Cookies, cookies_path: Path) -> None:
-    """Write cookies to file."""
-    cookies_dict = dict(cookies)
-
-    # Overwrite the cookies file
-    with cookies_path.open(mode="w") as f:
-        json.dump(cookies_dict, f)
-
-
-def read_cookies(cookies_path: Path) -> httpx.Cookies:
-    """Read cookies from file."""
-    try:
-        with cookies_path.open() as f:
-            cookies_dict = json.load(f)
-        return httpx.Cookies(cookies_dict)
-    except (FileNotFoundError, json.JSONDecodeError):
-        return httpx.Cookies()
-
-
-def delete_cookies(cookies_path: Path) -> None:
-    """Delete cookies file."""
-    cookies_path.unlink(missing_ok=True)
+def pprint_json(data: Any):
+    """Pretty print data."""
+    rich.print(orjson.dumps(data, option=orjson.OPT_INDENT_2).decode())