utils.py

import sys
import logging
import time
import base64
import jwt
import httpx

from cryptography.hazmat.primitives.serialization import load_der_private_key

logging.basicConfig(stream=sys.stdout, level=logging.INFO)
logger = logging.getLogger("Code Review Assistant")

BASE_GITHUB_URL = "https://api.github.com"

class JWTGenerator:
    def __init__(self, app_id, private_key):
        self.app_id = app_id
        self.private_key = private_key

    def generate_jwt(self):
        payload = {
            "iat": int(time.time()),
            "exp": int(time.time()) + (10 * 60),
            "iss": self.app_id,
        }
        if self.private_key:
            private_key_cleaned = self.private_key.replace("-----BEGIN RSA PRIVATE KEY-----", "").replace("\n", "").replace("-----END RSA PRIVATE KEY-----", "")
            secret = base64.b64decode(private_key_cleaned)

            private_rsa_key = load_der_private_key(secret, password=None)
            jwt_token = jwt.encode(payload, private_rsa_key, algorithm="RS256")
            return jwt_token
        raise ValueError("PRIVATE_KEY not found.")

async def get_installation_access_token(jwt, installation_id):
    url = f"{BASE_GITHUB_URL}/app/installations/{installation_id}/access_tokens"
    headers = {
        "Authorization": f"Bearer {jwt}",
        "Accept": "application/vnd.github.v3+json",
    }
    async with httpx.AsyncClient() as client:
        response = await client.post(url, headers=headers)
        return response.json()["token"]

def get_diff_url(pr):
    """GitHub 302s to this URL."""
    original_url = pr.get("url")
    parts = original_url.split("/")
    owner, repo, pr_number = parts[-4], parts[-3], parts[-1]
    return f"https://patch-diff.githubusercontent.com/raw/{owner}/{repo}/pull/{pr_number}.diff"
    # return f"{BASE_GITHUB_URL}/repos/{owner}/{repo}/pulls/{pr_number}"