Fix spec typo, wire Cognito groups to OPA, clean up stale files

spec/draft-ztxp-02.md

@@ -134,20 +134,22 @@ This API is intentionally compatible with the **OpenID AuthZEN “evaluate”**
 ---
 
 ## 8. Reference Implementation
-The reference Python toolkit **`ztxp-v0.2.py`** includes:
+The reference Python toolkit **`ztxpv0.2.py`** includes:
 
-- `sign` – generate signed TAMs from YAML or JSON  
-- `verify` – validate signature and timestamp  
-- `broker` – HTTP evaluation service  
-- `opa` – optional Open Policy Agent integration  
+- `sign` – generate signed TAMs from YAML or JSON
+- `verify` – validate signature and timestamp
+- `broker` – HTTP evaluation service
+- `opa` – optional Open Policy Agent integration
 
-**Repository:** [https://github.com/cliffbell/ZTXP](https://github.com/cliffbell/ZTXP)
+**Repository:** [https://github.com/cjb00/ZTXP](https://github.com/cjb00/ZTXP)
 
 **Example workflow**
 ```bash
-python ztxp-v0.2.py sign tam.yaml signed.json
-python ztxp-v0.2.py broker --port 8080
-curl -X POST -H "Content-Type: application/json"      --data @signed.json      http://127.0.0.1:8080/ztxp/evaluate
+python ztxpv0.2.py sign tam.yaml signed.json
+python ztxpv0.2.py broker --port 8080
+curl -X POST -H "Content-Type: application/json" \
+     --data @signed.json \
+     http://127.0.0.1:8080/ztxp/evaluate
 ```
 
 ---

ztxb-aws-lab/app/lambdas/pep_authorizer/handler.py

@@ -50,7 +50,16 @@ def build_tam(event):
     principal_id = "anonymous"
     groups = []
     if auth_header:
-        principal_id = _extract_sub_from_jwt(auth_header) or auth_header[:40]
+        claims = _decode_jwt_claims(auth_header)
+        principal_id = (
+            claims.get("sub")
+            or claims.get("email")
+            or claims.get("cognito:username")
+            or auth_header[:40]
+        )
+        # cognito:groups is a list of group names the user belongs to;
+        # these map directly to the OPA policy groups ("writer", "admin")
+        groups = claims.get("cognito:groups", [])
 
     # --- Device context (forwarded by client headers) ---
     device_id = headers.get("x-device-id", "unknown")
@@ -141,19 +150,22 @@ def call_broker(signed_tam):
 # JWT helper (lightweight — no external deps)
 # ---------------------------------------------------------------------------
 
-def _extract_sub_from_jwt(auth_header):
-    """Best-effort extraction of 'sub' claim from a Bearer JWT.
-    We do NOT verify the JWT here — Cognito + API Gateway handle that.
-    We only need the subject identifier for the TAM."""
+def _decode_jwt_claims(auth_header):
+    """Decode JWT claims from a Bearer token without verifying the signature.
+
+    Signature verification is handled upstream by Cognito + API Gateway.
+    We only need the claims (sub, cognito:groups, etc.) for the TAM.
+    Returns an empty dict on any parse failure.
+    """
     try:
         token = auth_header.replace("Bearer ", "").strip()
         payload_segment = token.split(".")[1]
+        # Restore base64 padding
         padding = 4 - len(payload_segment) % 4
         payload_segment += "=" * padding
-        claims = json.loads(base64.b64decode(payload_segment))
-        return claims.get("sub") or claims.get("email") or claims.get("cognito:username")
+        return json.loads(base64.b64decode(payload_segment))
     except Exception:
-        return None
+        return {}
 
 
 # ---------------------------------------------------------------------------

ztxb-aws-lab/infra/main.tf

@@ -21,8 +21,10 @@ module "kms" {
 ###############################################
 
 module "cognito" {
-  source  = "./modules/cognito"
-  project = var.project
+  source        = "./modules/cognito"
+  project       = var.project
+  callback_urls = var.cognito_callback_urls
+  logout_urls   = var.cognito_logout_urls
 }
 
 ###############################################

ztxb-aws-lab/infra/modules/cognito/main.tf

@@ -1,6 +1,6 @@
-variable "project" {
-  type = string
-}
+###############################################
+# USER POOL
+###############################################
 
 resource "aws_cognito_user_pool" "this" {
   name = "${var.project}-pool"
@@ -13,8 +13,34 @@ resource "aws_cognito_user_pool" "this" {
   }
 
   auto_verified_attributes = ["email"]
+
+  # cognito:groups is included automatically in the ID token
+  # when a user belongs to one or more groups
 }
 
+###############################################
+# USER POOL GROUPS
+# These map directly to the OPA policy groups:
+#   "writer" → can create/update notes
+#   "admin"  → full access
+###############################################
+
+resource "aws_cognito_user_group" "writer" {
+  user_pool_id = aws_cognito_user_pool.this.id
+  name         = "writer"
+  description  = "Users who can create and update notes"
+}
+
+resource "aws_cognito_user_group" "admin" {
+  user_pool_id = aws_cognito_user_pool.this.id
+  name         = "admin"
+  description  = "Administrators with full access"
+}
+
+###############################################
+# APP CLIENT
+###############################################
+
 resource "aws_cognito_user_pool_client" "app" {
   name         = "${var.project}-app-client"
   user_pool_id = aws_cognito_user_pool.this.id
@@ -25,15 +51,23 @@ resource "aws_cognito_user_pool_client" "app" {
   allowed_oauth_scopes                 = ["email", "openid", "profile"]
   allowed_oauth_flows_user_pool_client = true
 
-  callback_urls                = ["https://example.com/callback"]
-  logout_urls                  = ["https://example.com/logout"]
+  callback_urls                = var.callback_urls
+  logout_urls                  = var.logout_urls
   supported_identity_providers = ["COGNITO"]
 }
 
+###############################################
+# OUTPUTS
+###############################################
+
 output "user_pool_id" {
   value = aws_cognito_user_pool.this.id
 }
 
 output "user_pool_arn" {
   value = aws_cognito_user_pool.this.arn
 }
+
+output "app_client_id" {
+  value = aws_cognito_user_pool_client.app.id
+}

ztxb-aws-lab/infra/modules/cognito/variables.tf

@@ -0,0 +1,15 @@
+variable "project" {
+  type = string
+}
+
+variable "callback_urls" {
+  description = "OAuth callback URLs for the Cognito app client"
+  type        = list(string)
+  default     = ["https://example.com/callback"]
+}
+
+variable "logout_urls" {
+  description = "OAuth logout URLs for the Cognito app client"
+  type        = list(string)
+  default     = ["https://example.com/logout"]
+}

ztxb-aws-lab/infra/variables.tf

@@ -16,6 +16,18 @@ variable "pdp_image" {
   default     = "" # set via -var or tfvars (e.g. from build.sh output)
 }
 
+variable "cognito_callback_urls" {
+  description = "OAuth callback URLs for the Cognito app client"
+  type        = list(string)
+  default     = ["https://example.com/callback"]
+}
+
+variable "cognito_logout_urls" {
+  description = "OAuth logout URLs for the Cognito app client"
+  type        = list(string)
+  default     = ["https://example.com/logout"]
+}
+
 variable "tags" {
   description = "Default tags applied to resources"
   type        = map(string)

ztxb-aws-lab/tests/test_pep_authorizer.py

@@ -85,6 +85,38 @@ def test_jwt_subject_extraction(self):
         tam = pep.build_tam(event)
         assert tam["subject"]["id"] == "user:user-abc-123"
 
+    def test_cognito_groups_extracted(self):
+        claims = {"sub": "user-abc-123", "cognito:groups": ["writer", "admin"]}
+        payload = base64.b64encode(json.dumps(claims).encode()).decode().rstrip("=")
+        fake_jwt = f"eyJhbGciOiJSUzI1NiJ9.{payload}.fake-sig"
+
+        event = _make_event(auth_header=f"Bearer {fake_jwt}")
+        tam = pep.build_tam(event)
+        assert tam["subject"]["groups"] == ["writer", "admin"]
+
+    def test_no_groups_when_missing(self):
+        claims = {"sub": "user-abc-123"}
+        payload = base64.b64encode(json.dumps(claims).encode()).decode().rstrip("=")
+        fake_jwt = f"eyJhbGciOiJSUzI1NiJ9.{payload}.fake-sig"
+
+        event = _make_event(auth_header=f"Bearer {fake_jwt}")
+        tam = pep.build_tam(event)
+        assert tam["subject"]["groups"] == []
+
+
+class TestDecodeJwtClaims:
+    def test_decodes_claims(self):
+        claims = {"sub": "abc", "cognito:groups": ["writer"]}
+        payload = base64.b64encode(json.dumps(claims).encode()).decode().rstrip("=")
+        fake_jwt = f"header.{payload}.sig"
+        result = pep._decode_jwt_claims(f"Bearer {fake_jwt}")
+        assert result["sub"] == "abc"
+        assert result["cognito:groups"] == ["writer"]
+
+    def test_returns_empty_on_bad_token(self):
+        result = pep._decode_jwt_claims("Bearer not.valid")
+        assert result == {}
+
 
 class TestCanonicalJson:
     def test_deterministic(self):