refactor: Update discussion labeler and JWT generation to use environment variables for keys

Author: BethanyJep

.github/workflows/discussion-labeler.yml

@@ -27,10 +27,14 @@ jobs:
 
       - name: Run discussion labeler
         env:
-          TOKEN: ${{ secrets.GITHUB_TOKEN }}
+          TOKEN: ${{ secrets.TOKEN }}
           APP_ID: ${{ secrets.APP_ID }}
-          APP_PRIVATE_KEY_PATH: ${{ secrets.APP_PRIVATE_KEY_PATH }}
+          APP_PRIVATE_KEY: ${{ secrets.APP_PRIVATE_KEY }}
           APP_INSTALLATION_ID: ${{ secrets.APP_INSTALLATION_ID }}
+          AZURE_OPENAI_API_VERSION: ${{ secrets.AZURE_OPENAI_API_VERSION }}
+          AZURE_OPENAI_ENDPOINT: ${{ secrets.AZURE_OPENAI_ENDPOINT }}
+          AZURE_OPENAI_KEY: ${{ secrets.AZURE_OPENAI_KEY }}
+          SECRET_KEY: ${{ secrets.SECRET_KEY }}
           DEFAULT_REPO: ${{ github.repository }}
         run: |
-          python test.py ${{ secrets.SECRET_KEY }}
+          python basic.py

.github/workflows/on-discussion-created.yml

@@ -45,15 +45,24 @@
     json_tracer = PromptyTracer()
     Tracer.add("PromptyTracer", json_tracer.tracer)
 
-# API settings - get token from environment variables
+# Global variables
 TOKEN = os.getenv("TOKEN")
 DEFAULT_REPO = os.getenv("DEFAULT_REPO", "golclinics/discussions")
 REQUEST_TIMEOUT = int(os.getenv("REQUEST_TIMEOUT", "30"))  # Default 30 second timeout
 
 # App settings
 APP_ID = os.getenv("APP_ID")
-APP_PRIVATE_KEY_PATH = os.getenv("APP_PRIVATE_KEY_PATH", "./azure-ai-foundry-discussions.2025-05-06.private-key.pem")
-GITHUB_APP_INSTALLATION_ID = os.getenv("GITHUB_APP_INSTALLATION_ID")
+APP_PRIVATE_KEY = os.getenv("APP_PRIVATE_KEY")  # Direct key content instead of file path
+APP_PRIVATE_KEY_PATH = os.getenv("APP_PRIVATE_KEY_PATH")  # Keep for backward compatibility
+APP_INSTALLATION_ID = os.getenv("APP_INSTALLATION_ID")  # Removed GITHUB_ prefix
+
+# Azure OpenAI settings
+AZURE_OPENAI_API_VERSION = os.getenv("AZURE_OPENAI_API_VERSION")
+AZURE_OPENAI_ENDPOINT = os.getenv("AZURE_OPENAI_ENDPOINT")
+AZURE_OPENAI_KEY = os.getenv("AZURE_OPENAI_KEY")
+
+# Other settings
+SECRET_KEY = os.getenv("SECRET_KEY")
 
 # GitHub API rate limiting constants
 MAX_RETRIES = 3
@@ -69,21 +78,21 @@ class GithubAppAuthError(Exception):
 
 def validate_token() -> None:
     """Validate that the GitHub token is available."""
-    if not GITHUB_TOKEN:
+    if not TOKEN:
         raise TokenMissingError(
             "GitHub token not found in environment variables. "
-            "Please set the GITHUB_TOKEN environment variable."
+            "Please set the TOKEN environment variable."
         )
 
 def validate_github_app_config() -> None:
     """Validate that the GitHub App configuration is available."""
     missing = []
-    if not GITHUB_APP_ID:
-        missing.append("GITHUB_APP_ID")
-    if not GITHUB_APP_PRIVATE_KEY_PATH:
-        missing.append("GITHUB_APP_PRIVATE_KEY_PATH")
-    if not GITHUB_APP_INSTALLATION_ID:
-        missing.append("GITHUB_APP_INSTALLATION_ID")
+    if not APP_ID:
+        missing.append("APP_ID")
+    if not APP_PRIVATE_KEY and not APP_PRIVATE_KEY_PATH:
+        missing.append("APP_PRIVATE_KEY or APP_PRIVATE_KEY_PATH")
+    if not APP_INSTALLATION_ID:
+        missing.append("APP_INSTALLATION_ID")
 
     if missing:
         raise GithubAppAuthError(
@@ -111,22 +120,32 @@ def generate_jwt() -> str:
         payload = {
             'iat': now,               # Issued at time
             'exp': expiration,        # Expiration time
-            'iss': GITHUB_APP_ID      # GitHub App ID
+            'iss': APP_ID             # GitHub App ID
         }
 
-        # Read private key from file path
-        try:
-            key_path = Path(GITHUB_APP_PRIVATE_KEY_PATH)
-            if (key_path.exists()):
-                with open(key_path, "r") as key_file:
-                    private_key = key_file.read()
-                logger.info(f"Successfully loaded private key from file: {GITHUB_APP_PRIVATE_KEY_PATH}")
-            else:
-                logger.error(f"Private key file not found at: {GITHUB_APP_PRIVATE_KEY_PATH}")
-                raise FileNotFoundError(f"Private key file not found: {GITHUB_APP_PRIVATE_KEY_PATH}")
-        except Exception as e:
-            logger.error(f"Error reading private key file: {str(e)}")
-            raise
+        # Get the private key - either directly from env var or from file
+        private_key = None
+        
+        # First try to use the direct key from environment variable
+        if APP_PRIVATE_KEY:
+            private_key = APP_PRIVATE_KEY
+            logger.info("Using private key directly from environment variable")
+        # Fallback to loading from file path
+        elif APP_PRIVATE_KEY_PATH:
+            try:
+                key_path = Path(APP_PRIVATE_KEY_PATH)
+                if key_path.exists():
+                    with open(key_path, "r") as key_file:
+                        private_key = key_file.read()
+                    logger.info(f"Successfully loaded private key from file: {APP_PRIVATE_KEY_PATH}")
+                else:
+                    logger.error(f"Private key file not found at: {APP_PRIVATE_KEY_PATH}")
+                    raise FileNotFoundError(f"Private key file not found: {APP_PRIVATE_KEY_PATH}")
+            except Exception as e:
+                logger.error(f"Error reading private key file: {str(e)}")
+                raise
+        else:
+            raise GithubAppAuthError("No private key available. Set APP_PRIVATE_KEY or APP_PRIVATE_KEY_PATH")
 
         # Generate the JWT
         token = jwt.encode(
@@ -161,7 +180,7 @@ def get_installation_token() -> str:
         jwt_token = generate_jwt()
 
         # API endpoint for getting an installation token
-        url = f"https://api.github.com/app/installations/{GITHUB_APP_INSTALLATION_ID}/access_tokens"
+        url = f"https://api.github.com/app/installations/{APP_INSTALLATION_ID}/access_tokens"
 
         headers = {
             "Authorization": f"Bearer {jwt_token}",

basic.py

@@ -1,16 +1,97 @@
 import sys
+import os
+import logging
+import time
+import jwt
 
-APP_PRIVATE_KEY = None
-
-# get private key from args
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
+logger = logging.getLogger("test-jwt")
 
+# First try to get private key from args
+APP_PRIVATE_KEY = None
 if len(sys.argv) > 1:
     APP_PRIVATE_KEY = str(sys.argv[1])
 
+# If not provided as arg, try to get from environment variable
 if APP_PRIVATE_KEY is None:
-    print("Please provide the path to the private key file as an argument.")
-    sys.exit(1)
-else:
-    # Print only the first few characters as confirmation, not the entire key
-    masked_key = APP_PRIVATE_KEY[:5] + "..." if len(APP_PRIVATE_KEY) > 5 else "..."
-    print(f"Private key loaded successfully (begins with: {masked_key})")
+    APP_PRIVATE_KEY = os.getenv("APP_PRIVATE_KEY")
+
+# Also check if we have a file path available
+APP_PRIVATE_KEY_PATH = os.getenv("APP_PRIVATE_KEY_PATH")
+APP_ID = os.getenv("APP_ID")
+
+def test_jwt_generation():
+    """Test JWT generation with the private key"""
+    if not APP_ID:
+        logger.error("APP_ID environment variable is not set")
+        return False
+        
+    if not APP_PRIVATE_KEY and not APP_PRIVATE_KEY_PATH:
+        logger.error("Neither APP_PRIVATE_KEY nor APP_PRIVATE_KEY_PATH is available")
+        return False
+    
+    try:
+        # Use direct key if available
+        private_key = APP_PRIVATE_KEY
+        
+        # Otherwise try to load from file
+        if not private_key and APP_PRIVATE_KEY_PATH:
+            try:
+                with open(APP_PRIVATE_KEY_PATH, "r") as key_file:
+                    private_key = key_file.read()
+                logger.info(f"Successfully loaded private key from file: {APP_PRIVATE_KEY_PATH}")
+            except Exception as e:
+                logger.error(f"Error reading private key file: {str(e)}")
+                return False
+                
+        # JWT generation parameters
+        now = int(time.time())
+        expiration = now + (10 * 60)  # 10 minutes
+        
+        payload = {
+            'iat': now,               
+            'exp': expiration,        
+            'iss': APP_ID             
+        }
+        
+        # Generate the JWT
+        token = jwt.encode(
+            payload,
+            private_key,
+            algorithm='RS256'
+        )
+        
+        # If successful, print confirmation
+        logger.info("JWT generation successful!")
+        
+        # For security, only print the first few characters of the token
+        if isinstance(token, bytes):
+            token = token.decode('utf-8')
+        masked_token = token[:10] + "..." if len(token) > 10 else "..."
+        logger.info(f"Generated JWT token (begins with): {masked_token}")
+        
+        return True
+        
+    except Exception as e:
+        logger.error(f"JWT generation failed: {str(e)}")
+        return False
+
+if __name__ == "__main__":
+    if APP_PRIVATE_KEY is None and APP_PRIVATE_KEY_PATH is None:
+        logger.error("Please provide a private key as an argument or set APP_PRIVATE_KEY/APP_PRIVATE_KEY_PATH environment variable.")
+        sys.exit(1)
+    else:
+        # Mask key for security in logs
+        if APP_PRIVATE_KEY:
+            masked_key = APP_PRIVATE_KEY[:5] + "..." if len(APP_PRIVATE_KEY) > 5 else "..."
+            logger.info(f"Private key available (begins with: {masked_key})")
+        elif APP_PRIVATE_KEY_PATH:
+            logger.info(f"Private key path available: {APP_PRIVATE_KEY_PATH}")
+            
+        # Test JWT generation
+        if test_jwt_generation():
+            logger.info("JWT test successful")
+            sys.exit(0)
+        else:
+            logger.error("JWT test failed")
+            sys.exit(1)