From d8cbe6d73f18ebbcca29d7b83c4f0a47c6393fb1 Mon Sep 17 00:00:00 2001
From: meredith <meredith@sentry.io>
Date: Wed, 14 Dec 2022 16:04:21 -0800
Subject: [PATCH] feat(admin): implement basic admin user scopes

---
 snuba/admin/auth.py        | 11 +++++++++-
 snuba/admin/auth_scopes.py | 43 ++++++++++++++++++++++++++++++++++++++
 snuba/admin/user.py        |  6 +++++-
 3 files changed, 58 insertions(+), 2 deletions(-)
 create mode 100644 snuba/admin/auth_scopes.py

diff --git a/snuba/admin/auth.py b/snuba/admin/auth.py
index 394229eab8..c320a6f6ee 100644
--- a/snuba/admin/auth.py
+++ b/snuba/admin/auth.py
@@ -5,6 +5,7 @@
 from flask import request
 
 from snuba import settings
+from snuba.admin.auth_scopes import ADMIN_SCOPES
 from snuba.admin.jwt import validate_assertion
 from snuba.admin.user import AdminUser
 
@@ -28,7 +29,15 @@ def authorize_request() -> AdminUser:
     provider = AUTH_PROVIDERS.get(provider_id)
     if provider is None:
         raise ValueError("Invalid authorization provider")
-    return provider()
+
+    return _set_scopes(provider())
+
+
+def _set_scopes(user: AdminUser) -> AdminUser:
+    # todo: depending on provider convert user email
+    # to subset of ADMIN_SCOPES based on IAM roles
+    user.scopes = ADMIN_SCOPES
+    return user
 
 
 def passthrough_authorize() -> AdminUser:
diff --git a/snuba/admin/auth_scopes.py b/snuba/admin/auth_scopes.py
new file mode 100644
index 0000000000..5b699c1e7b
--- /dev/null
+++ b/snuba/admin/auth_scopes.py
@@ -0,0 +1,43 @@
+from dataclasses import dataclass
+from enum import Enum
+from typing import Set
+
+
+class AdminRole(Enum):
+    ADMIN = "admin"
+    MEMBER = "member"
+    MEMBER_READ = "member_read"
+
+
+ADMIN_CATEGORY_RESOURCES = {
+    # clickhouse migrations
+    "migrations.all",
+    "migrations.system",
+    "migrations.generic_metrics",
+    "migrations.profiles",
+    "migrations.functions",
+    "migrations.replays",
+    "migrations.querylog",
+    "migrations.test_migration",
+}
+
+
+@dataclass(frozen=True)
+class AuthScope:
+    category: str
+    resource: str
+    role: AdminRole
+
+    def to_str(self) -> str:
+        return f"{self.category}.{self.resource}.{self.role.value}"
+
+
+def scopes() -> Set[AuthScope]:
+    scopes = set()
+    for item in ADMIN_CATEGORY_RESOURCES:
+        category, resource = item.split(".")
+        scopes.update([AuthScope(category, resource, role) for role in AdminRole])
+    return scopes
+
+
+ADMIN_SCOPES = scopes()
diff --git a/snuba/admin/user.py b/snuba/admin/user.py
index 91b60a4b81..728ab62c0b 100644
--- a/snuba/admin/user.py
+++ b/snuba/admin/user.py
@@ -1,4 +1,7 @@
-from dataclasses import dataclass
+from dataclasses import dataclass, field
+from typing import Set
+
+from snuba.admin.auth_scopes import AuthScope
 
 
 @dataclass
@@ -10,3 +13,4 @@ class AdminUser:
 
     email: str
     id: str
+    scopes: Set[AuthScope] = field(default_factory=set)