From 5079daa4ad837cb5336c7f8a6de98e50d43a3c24 Mon Sep 17 00:00:00 2001 From: Ori Kotek Date: Thu, 24 Aug 2023 16:33:51 +0300 Subject: [PATCH] Bitbucket server, WIP --- pr_agent/git_providers/bitbucket_provider.py | 7 +- pr_agent/secret_providers/__init__.py | 16 ++++ .../google_cloud_storage_secret_provider.py | 35 +++++++++ pr_agent/secret_providers/secret_provider.py | 12 +++ pr_agent/servers/bitbucket_app.py | 74 ++++++++++++++++++- pr_agent/settings/configuration.toml | 1 + pyproject.toml | 4 +- requirements.txt | 2 + 8 files changed, 145 insertions(+), 6 deletions(-) create mode 100644 pr_agent/secret_providers/__init__.py create mode 100644 pr_agent/secret_providers/google_cloud_storage_secret_provider.py create mode 100644 pr_agent/secret_providers/secret_provider.py diff --git a/pr_agent/git_providers/bitbucket_provider.py b/pr_agent/git_providers/bitbucket_provider.py index bee0e351..6da864e4 100644 --- a/pr_agent/git_providers/bitbucket_provider.py +++ b/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) diff --git a/pr_agent/secret_providers/__init__.py b/pr_agent/secret_providers/__init__.py new file mode 100644 index 00000000..1cc3ea7b --- /dev/null +++ b/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 diff --git a/pr_agent/secret_providers/google_cloud_storage_secret_provider.py b/pr_agent/secret_providers/google_cloud_storage_secret_provider.py new file mode 100644 index 00000000..18db5c4b --- /dev/null +++ b/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 diff --git a/pr_agent/secret_providers/secret_provider.py b/pr_agent/secret_providers/secret_provider.py new file mode 100644 index 00000000..df1e7780 --- /dev/null +++ b/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 diff --git a/pr_agent/servers/bitbucket_app.py b/pr_agent/servers/bitbucket_app.py index 3b21ca7a..039fc5c8 100644 --- a/pr_agent/servers/bitbucket_app.py +++ b/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) diff --git a/pr_agent/settings/configuration.toml b/pr_agent/settings/configuration.toml index c2350526..f8abd555 100644 --- a/pr_agent/settings/configuration.toml +++ b/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 diff --git a/pyproject.toml b/pyproject.toml index 9a945dca..8d429668 100644 --- a/pyproject.toml +++ b/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] diff --git a/requirements.txt b/requirements.txt index 0bbb6f28..fe92a74b 100644 --- a/requirements.txt +++ b/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