litellm-mirror/litellm/llms/github_copilot/authenticator.py

import os

import httpx
import json

import time
from datetime import datetime
from typing import Optional

from litellm._logging import verbose_logger
from litellm.caching import InMemoryCache
from litellm.llms.custom_httpx.http_handler import (
    _get_httpx_client,
    get_async_httpx_client,
    httpxSpecialProvider,
)

from .constants import (
    GetAccessTokenError,
    GetDeviceCodeError,
    RefreshAPIKeyError,
    GetAPIKeyError,
    APIKeyExpiredError,
)


class Authenticator:
    def __init__(self) -> None:
        # Token storage paths
        self.token_dir = os.getenv(
            "GITHUB_COPILOT_TOKEN_DIR",
            os.path.expanduser("~/.config/litellm/github_copilot"),
        )
        self.access_token_file = os.path.join(
            self.token_dir,
            os.getenv("GITHUB_COPILOT_ACCESS_TOKEN_FILE", "access-token"),
        )
        self.api_key_file = os.path.join(
            self.token_dir, os.getenv("GITHUB_COPILOT_API_KEY_FILE", "api-key.json")
        )
        self._ensure_token_dir()

    def get_access_token(self) -> str:
        """
        Login to Copilot with retry 3 times

        Returns:
        access_token: str

        """
        try:
            with open(self.access_token_file, "r") as f:
                access_token = f.read().strip()
                return access_token
        except IOError:
            verbose_logger.warning("Error loading access token from file")

        for _ in range(3):
            try:
                access_token = self._login()
            except GetDeviceCodeError | GetAccessTokenError | RefreshAPIKeyError:
                continue
            else:
                try:
                    with open(self.access_token_file, "w") as f:
                        f.write(access_token)
                except IOError:
                    verbose_logger.error("Error saving access token to file")
                return access_token

        raise GetAccessTokenError("Failed to get access token")

    def get_api_key(self) -> str:
        """Get the API key"""
        try:
            with open(self.api_key_file, "r") as f:
                api_key_info = json.load(f)
                if api_key_info.get("expires_at") > datetime.now().timestamp():
                    return api_key_info.get("token")
                else:
                    raise APIKeyExpiredError("API key expired")
        except IOError:
            verbose_logger.warning("Error opening API key file")
        except (json.JSONDecodeError, KeyError):
            verbose_logger.warning("Error reading API key from file")
        except APIKeyExpiredError:
            verbose_logger.warning("API key expired")

        try:
            api_key_info = self._refresh_api_key()
            with open(self.api_key_file, "w") as f:
                json.dump(api_key_info, f)
        except IOError:
            verbose_logger.error("Error saving API key to file")
        except RefreshAPIKeyError:
            raise GetAPIKeyError("Failed to refresh API key")

        return api_key_info.get("token")

    def _refresh_api_key(self) -> dict:
        """
        Refresh the API key using the access token

        Returns:
        api_key_info: dict
        """

        access_token = self.get_access_token()
        headers = {
            "authorization": f"token {access_token}",
            "editor-version": "vscode/1.85.1",
            "editor-plugin-version": "copilot/1.155.0",
            "user-agent": "GithubCopilot/1.155.0",
        }

        max_retries = 3
        for _ in range(max_retries):
            try:
                sync_client = _get_httpx_client()
                response = sync_client.get(
                    "https://api.github.com/copilot_internal/v2/token", headers=headers
                )
                response.raise_for_status()

                response_json = response.json()

                if "token" in response_json:
                    return response_json
            except httpx.HTTPStatusError as e:
                verbose_logger.error(f"Error refreshing API key: {str(e)}")

        raise RefreshAPIKeyError("Failed to refresh API key")

    def _ensure_token_dir(self) -> None:
        """Ensure the token directory exists"""
        if not os.path.exists(self.token_dir):
            os.makedirs(self.token_dir, exist_ok=True)

    def _login(self) -> str:
        """
        Login to GitHub Copilot using device code flow

        Returns:
        access_token: str
        """

        try:
            sync_client = _get_httpx_client()
            # Get device code
            resp = sync_client.post(
                "https://github.com/login/device/code",
                headers={
                    "accept": "application/json",
                    "editor-version": "vscode/1.85.1",
                    "editor-plugin-version": "copilot/1.155.0",
                    "content-type": "application/json",
                    "user-agent": "GithubCopilot/1.155.0",
                    "accept-encoding": "gzip,deflate,br",
                },
                json={"client_id": "Iv1.b507a08c87ecfe98", "scope": "read:user"},
            )
            resp.raise_for_status()
            resp_json = resp.json()

            device_code = resp_json.get("device_code")
            user_code = resp_json.get("user_code")
            verification_uri = resp_json.get("verification_uri")

            if not all([device_code, user_code, verification_uri]):
                verbose_logger.error("Response missing required fields")
                return None
        except httpx.HttpError as e:
            verbose_logger.error(f"Error getting device code: {str(e)}")
            raise GetDeviceCodeError("Failed to get device code")
        except json.JSONDecodeError as e:
            verbose_logger.error(f"Error decoding JSON response: {str(e)}")
            raise GetDeviceCodeError("Failed to get device code")
        except RuntimeError as e:
            verbose_logger.error(f"Error getting device code: {str(e)}")
            raise GetDeviceCodeError("Failed to get device code")

        print(
            f"Please visit {verification_uri} and enter code {user_code} to authenticate."
        )

        while True:
            time.sleep(5)

            # Get access token
            try:
                resp = sync_client.post(
                    "https://github.com/login/oauth/access_token",
                    headers={
                        "accept": "application/json",
                        "editor-version": "vscode/1.85.1",
                        "editor-plugin-version": "copilot/1.155.0",
                        "content-type": "application/json",
                        "user-agent": "GithubCopilot/1.155.0",
                        "accept-encoding": "gzip,deflate,br",
                    },
                    json={
                        "client_id": "Iv1.b507a08c87ecfe98",
                        "device_code": device_code,
                        "grant_type": "urn:ietf:params:oauth:grant-type:device_code",
                    },
                )
                resp.raise_for_status()
                resp_json = resp.json()

                if "access_token" in resp_json:
                    verbose_logger.info("Authentication success!")
                    return resp_json["access_token"]
                else:
                    continue
            except httpx.HTTPStatusError as e:
                verbose_logger.error(f"Error getting access token: {str(e)}")
                raise GetAccessTokenError("Failed to get access token")
            except json.JSONDecodeError as e:
                verbose_logger.error(f"Error decoding JSON response: {str(e)}")
                raise GetAccessTokenError("Failed to get access token")