feat: Add managed_secrets app and integrate Azure Key Vault

daras_ai_v2/manage_api_keys_widget.py

@@ -1,6 +1,8 @@
 import typing
 
 import gooey_gui as gui
+from django.contrib.humanize.templatetags.humanize import naturaltime
+
 from api_keys.models import ApiKey
 from app_users.models import AppUser
 from daras_ai_v2 import icons
@@ -25,41 +27,69 @@ def manage_api_keys(workspace: "Workspace", user: AppUser):
 
     api_keys = list(workspace.api_keys.order_by("-created_at"))
 
-    table_area = gui.div()
+    table_area = gui.div(
+        className="table-responsive text-nowrap container-margin-reset"
+    )
 
-    if gui.button("＋ Create new secret key"):
+    if gui.button("＋ Create new API Key"):
         api_key = generate_new_api_key(workspace=workspace, user=user)
         api_keys.insert(0, api_key)
 
     with table_area:
-        import pandas as pd
-
-        gui.table(
-            pd.DataFrame.from_records(
-                columns=["Secret Key (Preview)", "Created At", "Created By"],
-                data=[
-                    (
-                        api_key.preview,
-                        api_key.created_at.strftime("%B %d, %Y at %I:%M:%S %p %Z"),
-                        api_key.created_by and api_key.created_by.full_name() or "",
-                    )
-                    for api_key in api_keys
-                ],
-            ),
-        )
+        with gui.tag("table", className="table table-striped"):
+            with gui.tag("thead"), gui.tag("tr"):
+                with gui.tag("th"):
+                    gui.write("Gooey.AI Key")
+                with gui.tag("th"):
+                    gui.write("Created At")
+                with gui.tag("th"):
+                    gui.write("Created By")
+                gui.tag("th")
+            with gui.tag("tbody"):
+                for api_key in api_keys:
+                    with gui.tag("tr"):
+                        with gui.tag("td"):
+                            gui.write(f"`{api_key.preview}`")
+                        with gui.tag("td"):
+                            gui.write(str(naturaltime(api_key.created_at)))
+                        with gui.tag("td"):
+                            gui.write(
+                                api_key.created_by
+                                and api_key.created_by.full_name()
+                                or ""
+                            )
+                        with gui.tag("td"):
+                            delete_dialog = gui.use_confirm_dialog(
+                                key=f"delete_api_key_{api_key.id}"
+                            )
+                            gui.button_with_confirm_dialog(
+                                ref=delete_dialog,
+                                trigger_label=icons.delete,
+                                trigger_type="tertiary",
+                                trigger_className="text-danger p-0 m-0",
+                                modal_title="### Delete API Key",
+                                modal_content=f"Are you sure you want to delete `{api_key.preview}`?\n\n"
+                                "API requests made using this key will be rejected, "
+                                "which could cause any systems still depending on it to break. "
+                                "Once deleted, you'll no longer be able to view or modify this API key.",
+                                confirm_label="Delete",
+                                confirm_className="border-danger bg-danger text-white",
+                            )
+                            if delete_dialog.pressed_confirm:
+                                api_key.delete()
+                                gui.rerun()
 
 
 def generate_new_api_key(workspace: "Workspace", user: AppUser) -> ApiKey:
     api_key, secret_key = ApiKey.objects.create_api_key(workspace, created_by=user)
 
     gui.success(
         """
-##### API key generated
-
+**API key generated**
 Please save this secret key somewhere safe and accessible.
 For security reasons, **you won't be able to view it again** through your account.
 If you lose this secret key, you'll need to generate a new one.
-            """
+        """
     )
     col1, col2 = gui.columns([3, 1], responsive=False)
     with col1:

daras_ai_v2/settings.py

@@ -73,6 +73,7 @@
     "functions",
     "workspaces",
     "api_keys",
+    "managed_secrets",
 ]
 
 MIDDLEWARE = [
@@ -424,6 +425,8 @@
 AZURE_SPEECH_ENDPOINT = f"https://{AZURE_SPEECH_REGION}.api.cognitive.microsoft.com"
 AZURE_TTS_ENDPOINT = f"https://{AZURE_SPEECH_REGION}.tts.speech.microsoft.com"
 
+AZURE_KEY_VAULT_ENDPOINT = config("AZURE_KEY_VAULT_ENDPOINT", "")
+
 AZURE_OPENAI_ENDPOINT_CA = config("AZURE_OPENAI_ENDPOINT_CA", "")
 AZURE_OPENAI_KEY_CA = config("AZURE_OPENAI_KEY_CA", "")
 AZURE_OPENAI_ENDPOINT_EASTUS2 = config("AZURE_OPENAI_ENDPOINT_EASTUS2", "")

functions/executor.js

@@ -8,12 +8,17 @@ Deno.serve(async (req) => {
     return new Response("Unauthorized", { status: 401 });
   }
 
-  let { tag, code, variables } = await req.json();
+  let { tag, code, variables, env } = await req.json();
   let { mockConsole, logs } = captureConsole(tag);
   let status, response;
 
   try {
-    let retval = isolatedEval(mockConsole, code, variables);
+    let retval = isolatedEval({
+      console: mockConsole,
+      code,
+      variables,
+      process: { env },
+    });
     if (retval instanceof Function) {
       retval = retval(variables);
     }
@@ -31,7 +36,7 @@ Deno.serve(async (req) => {
   return new Response(body, { status });
 });
 
-function isolatedEval(console, code, variables) {
+function isolatedEval({ console, code, variables, process }) {
   // Hide global objects
   let Deno, global, self, globalThis, window;
   return eval(code);

managed_secrets/admin.py

@@ -0,0 +1,18 @@
+from django.contrib import admin
+
+from managed_secrets.models import ManagedSecret
+
+
+@admin.register(ManagedSecret)
+class ManagedSecretAdmin(admin.ModelAdmin):
+    list_display = [
+        "__str__",
+        "workspace",
+        "created_at",
+        "updated_at",
+        "usage_count",
+        "last_used_at",
+    ]
+    autocomplete_fields = ["workspace", "created_by"]
+    readonly_fields = ["usage_count", "last_used_at", "external_id"]
+    search_fields = ["name", "external_id"]

managed_secrets/apps.py

@@ -0,0 +1,11 @@
+from django.apps import AppConfig
+
+
+class ManagedSecretsConfig(AppConfig):
+    default_auto_field = "django.db.models.BigAutoField"
+    name = "managed_secrets"
+
+    def ready(self):
+        from . import signals
+
+        assert signals

managed_secrets/migrations/0001_initial.py

@@ -0,0 +1,35 @@
+# Generated by Django 5.1.3 on 2024-12-23 15:47
+
+import django.db.models.deletion
+import uuid
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    initial = True
+
+    dependencies = [
+        ('app_users', '0023_alter_appusertransaction_workspace'),
+        ('workspaces', '0006_workspace_description'),
+    ]
+
+    operations = [
+        migrations.CreateModel(
+            name='ManagedSecret',
+            fields=[
+                ('id', models.BigAutoField(auto_created=True, primary_key=True, serialize=False, verbose_name='ID')),
+                ('external_id', models.UUIDField(default=uuid.uuid4, editable=False)),
+                ('name', models.CharField(max_length=255)),
+                ('created_at', models.DateTimeField(auto_now_add=True)),
+                ('updated_at', models.DateTimeField(auto_now=True)),
+                ('usage_count', models.PositiveIntegerField(db_index=True, default=0)),
+                ('last_used_at', models.DateTimeField(blank=True, default=None, null=True)),
+                ('created_by', models.ForeignKey(null=True, on_delete=django.db.models.deletion.SET_NULL, to='app_users.appuser')),
+                ('workspace', models.ForeignKey(on_delete=django.db.models.deletion.CASCADE, related_name='managed_secrets', to='workspaces.workspace')),
+            ],
+            options={
+                'unique_together': {('workspace', 'name')},
+            },
+        ),
+    ]

managed_secrets/models.py

@@ -0,0 +1,74 @@
+import uuid
+
+from django.db import models
+
+from daras_ai_v2 import settings
+
+'az role assignment create --role "Key Vault Secrets Officer" --assignee "<upn>" --scope "/subscriptions/<subscription-id>/resourceGroups/<resource-group-name>/providers/Microsoft.KeyVault/vaults/<your-unique-keyvault-name>"'
+
+
+class ManagedSecretQuerySet(models.QuerySet):
+    def create(self, *, value: str, **kwargs):
+        secret = super().create(**kwargs)
+        secret.store_value(value)
+        return secret
+
+
+class ManagedSecret(models.Model):
+    external_id = models.UUIDField(default=uuid.uuid4, editable=False)
+
+    name = models.CharField(max_length=255)
+
+    workspace = models.ForeignKey(
+        "workspaces.Workspace", on_delete=models.CASCADE, related_name="managed_secrets"
+    )
+    created_by = models.ForeignKey(
+        "app_users.AppUser", on_delete=models.SET_NULL, null=True
+    )
+
+    created_at = models.DateTimeField(auto_now_add=True)
+    updated_at = models.DateTimeField(auto_now=True)
+
+    usage_count = models.PositiveIntegerField(default=0, db_index=True)
+    last_used_at = models.DateTimeField(null=True, blank=True, default=None)
+
+    objects: ManagedSecretQuerySet = ManagedSecretQuerySet.as_manager()
+
+    value: str | None = None
+
+    class Meta:
+        unique_together = ("workspace", "name")
+
+    def __str__(self):
+        return f"{self.name} ({self._external_name()})"
+
+    def store_value(self, value: str):
+        client = _get_az_secret_client()
+        client.set_secret(self._external_name(), value)
+        self.value = value
+
+    def load_value(self):
+        client = _get_az_secret_client()
+        self.value = client.get_secret(self._external_name()).value
+
+    def delete_value(self):
+        import azure.core.exceptions
+
+        client = _get_az_secret_client()
+        try:
+            client.begin_delete_secret(self._external_name())
+        except azure.core.exceptions.ResourceNotFoundError:
+            pass
+
+    def _external_name(self):
+        return f"ms-gooey-{self.external_id}"
+
+
+def _get_az_secret_client():
+    from azure.keyvault.secrets import SecretClient
+    from azure.identity import DefaultAzureCredential
+
+    credential = DefaultAzureCredential()
+    return SecretClient(
+        vault_url=settings.AZURE_KEY_VAULT_ENDPOINT, credential=credential
+    )

managed_secrets/signals.py

@@ -0,0 +1,16 @@
+from django.db.models.signals import post_save, post_delete
+from django.dispatch import receiver
+
+from managed_secrets.models import ManagedSecret
+
+
+@receiver(post_save, sender=ManagedSecret)
+def on_secret_saved(instance: ManagedSecret, **kwargs):
+    if instance.value is None:
+        return
+    instance.store_value(instance.value)
+
+
+@receiver(post_delete, sender=ManagedSecret)
+def on_secret_deleted(instance: ManagedSecret, **kwargs):
+    instance.delete_value()

managed_secrets/tests.py

@@ -0,0 +1,3 @@
+from django.test import TestCase
+
+# Create your tests here.