Bitbucket server, WIP

pr_agent/git_providers/bitbucket_provider.py

@@ -5,6 +5,7 @@ from urllib.parse import urlparse
 
 import requests
 from atlassian.bitbucket import Cloud
+from starlette_context import context
 
 from ..config_loader import get_settings
 from .git_provider import FilePatchInfo, GitProvider
@@ -13,7 +14,11 @@ from .git_provider import FilePatchInfo, GitProvider
 class BitbucketProvider(GitProvider):
     def __init__(self, pr_url: Optional[str] = None, incremental: Optional[bool] = False):
         s = requests.Session()
-        s.headers['Authorization'] = f'Bearer {get_settings().get("BITBUCKET.BEARER_TOKEN", None)}'
+        try:
+            bearer = context.get("bitbucket_bearer_token", None)
+            s.headers['Authorization'] = f'Bearer {bearer}'
+        except Exception:
+            s.headers['Authorization'] = f'Bearer {get_settings().get("BITBUCKET.BEARER_TOKEN", None)}'
         s.headers['Content-Type'] = 'application/json'
         self.headers = s.headers
         self.bitbucket_client = Cloud(session=s)

pr_agent/secret_providers/__init__.py

@@ -0,0 +1,16 @@
+from pr_agent.config_loader import get_settings
+
+
+def get_secret_provider():
+    try:
+        provider_id = get_settings().config.secret_provider
+    except AttributeError as e:
+        raise ValueError("secret_provider is a required attribute in the configuration file") from e
+    try:
+        if provider_id == 'google_cloud_storage':
+            from pr_agent.secret_providers.google_cloud_storage_secret_provider import GoogleCloudStorageSecretProvider
+            return GoogleCloudStorageSecretProvider()
+        else:
+            raise ValueError(f"Unknown secret provider: {provider_id}")
+    except Exception as e:
+        raise ValueError(f"Failed to initialize secret provider {provider_id}") from e

pr_agent/secret_providers/google_cloud_storage_secret_provider.py

@@ -0,0 +1,35 @@
+import ujson
+
+from google.cloud import storage
+
+from pr_agent.config_loader import get_settings
+from pr_agent.git_providers.gitlab_provider import logger
+from pr_agent.secret_providers.secret_provider import SecretProvider
+
+
+class GoogleCloudStorageSecretProvider(SecretProvider):
+    def __init__(self):
+        try:
+            self.client = storage.Client.from_service_account_info(ujson.loads(get_settings().google_cloud_storage.
+                                                                               service_account))
+            self.bucket_name = get_settings().google_cloud_storage.bucket_name
+            self.bucket = self.client.bucket(self.bucket_name)
+        except Exception as e:
+            logger.error(f"Failed to initialize Google Cloud Storage Secret Provider: {e}")
+            raise e
+
+    def get_secret(self, secret_name: str) -> str:
+        try:
+            blob = self.bucket.blob(secret_name)
+            return blob.download_as_string()
+        except Exception as e:
+            logger.error(f"Failed to get secret {secret_name} from Google Cloud Storage: {e}")
+            return ""
+
+    def store_secret(self, secret_name: str, secret_value: str):
+        try:
+            blob = self.bucket.blob(secret_name)
+            blob.upload_from_string(secret_value)
+        except Exception as e:
+            logger.error(f"Failed to store secret {secret_name} in Google Cloud Storage: {e}")
+            raise e

pr_agent/secret_providers/secret_provider.py

@@ -0,0 +1,12 @@
+from abc import ABC, abstractmethod
+
+
+class SecretProvider(ABC):
+
+    @abstractmethod
+    def get_secret(self, secret_name: str) -> str:
+        pass
+
+    @abstractmethod
+    def store_secret(self, secret_name: str, secret_value: str):
+        pass

pr_agent/servers/bitbucket_app.py

@@ -1,16 +1,52 @@
+import hashlib
 import json
+import logging
 import os
+import time
 
+import jwt
+import requests
 import uvicorn
 from fastapi import APIRouter, FastAPI, Request, Response
 from starlette.middleware import Middleware
 from starlette.responses import JSONResponse
+from starlette_context import context
 from starlette_context.middleware import RawContextMiddleware
 
+from pr_agent.agent.pr_agent import PRAgent
 from pr_agent.config_loader import get_settings
+from pr_agent.secret_providers import get_secret_provider
 
 router = APIRouter()
+secret_provider = get_secret_provider()
 
+async def get_bearer_token(shared_secret: str, client_key: str):
+    try:
+        now = int(time.time())
+        url = "https://bitbucket.org/site/oauth2/access_token"
+        canonical_url = "GET&/site/oauth2/access_token&"
+        qsh = hashlib.sha256(canonical_url.encode("utf-8")).hexdigest()
+        app_key = get_settings().bitbucket.app_key
+
+        payload = {
+            "iss": app_key,
+            "iat": now,
+            "exp": now + 240,
+            "qsh": qsh,
+            "sub": client_key,
+            }
+        token = jwt.encode(payload, shared_secret, algorithm="HS256")
+        payload = 'grant_type=urn%3Abitbucket%3Aoauth2%3Ajwt'
+        headers = {
+            'Authorization': f'JWT {token}',
+            'Content-Type': 'application/x-www-form-urlencoded'
+        }
+        response = requests.request("POST", url, headers=headers, data=payload)
+        bearer_token = response.json()["access_token"]
+        return bearer_token
+    except Exception as e:
+        logging.error(f"Failed to get bearer token: {e}")
+        raise e
 
 @router.get("/")
 async def handle_manifest(request: Request, response: Response):
@@ -20,8 +56,24 @@ async def handle_manifest(request: Request, response: Response):
 
 @router.post("/webhook")
 async def handle_github_webhooks(request: Request, response: Response):
-    data = await request.json()
-    print(data)
+    try:
+        print(request.headers)
+        data = await request.json()
+        print(data)
+        owner = data["data"]["repository"]["owner"]["username"]
+        secrets = json.loads(secret_provider.get_secret(owner))
+        shared_secret = secrets["shared_secret"]
+        client_key = secrets["client_key"]
+        bearer_token = await get_bearer_token(shared_secret, client_key)
+        context['bitbucket_bearer_token'] = bearer_token
+        event = data["event"]
+        agent = PRAgent()
+        if event == "pullrequest:created":
+            pr_url = data["data"]["pullrequest"]["links"]["html"]["href"]
+            await agent.handle_request(pr_url, "review")
+    except Exception as e:
+        logging.error(f"Failed to handle webhook: {e}")
+        return JSONResponse({"error": "Unable to handle webhook"}, status_code=500)
 
 @router.get("/webhook")
 async def handle_github_webhooks(request: Request, response: Response):
@@ -29,8 +81,21 @@ async def handle_github_webhooks(request: Request, response: Response):
 
 @router.post("/installed")
 async def handle_installed_webhooks(request: Request, response: Response):
-    data = await request.json()
-    print(data)
+    try:
+        print(request.headers)
+        data = await request.json()
+        print(data)
+        shared_secret = data["sharedSecret"]
+        client_key = data["clientKey"]
+        username = data["principal"]["username"]
+        secrets = {
+            "shared_secret": shared_secret,
+            "client_key": client_key
+        }
+        secret_provider.store_secret(username, json.dumps(secrets))
+    except Exception as e:
+        logging.error(f"Failed to register user: {e}")
+        return JSONResponse({"error": "Unable to register user"}, status_code=500)
 
 @router.post("/uninstalled")
 async def handle_uninstalled_webhooks(request: Request, response: Response):
@@ -40,6 +105,7 @@ async def handle_uninstalled_webhooks(request: Request, response: Response):
 
 def start():
     get_settings().set("CONFIG.PUBLISH_OUTPUT_PROGRESS", False)
+    get_settings().set("CONFIG.GIT_PROVIDER", "bitbucket")
     middleware = [Middleware(RawContextMiddleware)]
     app = FastAPI(middleware=middleware)
     app.include_router(router)

pr_agent/settings/configuration.toml

@@ -11,6 +11,7 @@ ai_timeout=180
 max_description_tokens = 500
 max_commits_tokens = 500
 litellm_debugger=false
+secret_provider="google_cloud_storage"
 
 [pr_reviewer] # /review #
 require_focused_review=false

pyproject.toml

@@ -44,7 +44,9 @@ dependencies = [
   "starlette-context==0.3.6",
   "litellm~=0.1.445",
   "PyYAML==6.0",
-  "boto3~=1.28.25"
+  "boto3~=1.28.25",
+  "google-cloud-storage==2.10.0",
+  "ujson==5.8.0"
 ]
 
 [project.urls]

requirements.txt

@@ -15,3 +15,5 @@ PyYAML==6.0
 starlette-context==0.3.6
 litellm~=0.1.445
 boto3~=1.28.25
+google-cloud-storage==2.10.0
+ujson==5.8.0