Source code for util.github

from airflow.hooks.base import BaseHook

from config import Config
from util.base import BaseUtil




class GitHubUtil:
    """
    GitHub utility functions to interact with GitHub API
    """



    @staticmethod
    def get_github_creds() -> dict:
        """
        Retrieve GitHub App credentials from AWS Secrets Manager
        Parse fields for custom attributes

        :return: GitHub App credentials
        :rtype: dict
        """
        from json import loads

        github_creds = BaseHook.get_connection(Config.AWS_GITHUB_APP_TOKEN)
        # Parse the extra field which should contain our custom attributes
        extra = loads(github_creds.extra) if github_creds.extra else {}
        app_id = extra.get("app_id")
        installation_id = extra.get("installation_id")
        private_key = extra["private_key"].replace("\\n", "\n")

        return (
            {
                "app_id": app_id,
                "installation_id": installation_id,
                "private_key": private_key,
            }
            if all(
                [
                    app_id,
                    installation_id,
                    private_key,
                ]
            )
            else {}
        )




    @staticmethod
    def get_access_token(
        app_id: str,
        private_key: str,
        installation_id: str,
    ) -> str:
        """
        Get installation access token for GitHub App

        :param app_id: GitHub App ID
        :type app_id: str
        :param private_key: GitHub App private key
        :type private_key: str
        :param installation_id: GitHub App installation ID
        :type installation_id: str

        :return: Installation access token
        :rtype: str
        """
        import jwt
        import time
        import requests

        now = int(time.time())
        payload = {
            "iat": now,
            "exp": now + 600,
            "iss": app_id,
        }
        jwt_token = jwt.encode(
            payload,
            private_key,
            algorithm="RS256",
        )
        # Get installation access token
        headers = {
            "Authorization": f"Bearer {jwt_token}",
            "Accept": "application/vnd.github.v3+json",
        }
        # Use requests directly instead of BaseUtil.post_request
        try:
            response = requests.post(
                url=f"https://api.github.com/app/installations/{installation_id}/access_tokens",
                headers=headers,
            )
            if response.status_code == 201:
                return response.json()["token"]
            else:
                print(
                    f"Failed to get installation access token: {response.status_code} - {response.text}"
                )
                return None
        except Exception as e:
            print(f"Error getting GitHub access token: {str(e)}")
            return None