#!/usr/bin/env python3 import logging import random from dataclasses import dataclass from typing import Any, Dict, List, Optional, Union import boto3 # type: ignore from github import Github from github.AuthenticatedUser import AuthenticatedUser from github.GithubException import BadCredentialsException from github.NamedUser import NamedUser @dataclass class Token: user: Union[AuthenticatedUser, NamedUser] value: str rest: int SAFE_REQUESTS_LIMIT = 1000 def get_parameter_from_ssm( name: str, decrypt: bool = True, client: Optional[Any] = None ) -> str: if not client: client = boto3.client("ssm", region_name="us-east-1") return client.get_parameter( # type:ignore Name=name, WithDecryption=decrypt )[ "Parameter" ]["Value"] def get_parameters_from_ssm( names: List[str], decrypt: bool = True, client: Optional[Any] = None ) -> Dict[str, str]: if not client: client = boto3.client("ssm", region_name="us-east-1") names = list(set(names)) results = {} # type: Dict[str,str] i = 0 while (i) * 10 < len(names): # the get_parameters returns up to 10 values, so the call is split by 10 results.update( **{ p["Name"]: p["Value"] for p in client.get_parameters( Names=names[i * 10 : (i + 1) * 10], WithDecryption=decrypt )["Parameters"] } ) i += 1 return results ROBOT_TOKEN = None # type: Optional[Token] def get_best_robot_token(tokens_path: str = "/github-tokens") -> str: global ROBOT_TOKEN if ROBOT_TOKEN is not None: return ROBOT_TOKEN.value client = boto3.client("ssm", region_name="us-east-1") tokens = { p["Name"]: p["Value"] for p in client.get_parameters_by_path(Path=tokens_path, WithDecryption=True)[ "Parameters" ] } assert tokens token_items = list(tokens.items()) random.shuffle(token_items) best_token: Optional[Token] = None for name, value in token_items: gh = Github(value, per_page=100) try: # Do not spend additional request to API by accessing user.login unless # the token is chosen by the remaining requests number user = gh.get_user() rest, _ = gh.rate_limiting except BadCredentialsException: logging.error( "The token %(name)s has expired, please update it\n" "::error::Token %(name)s has expired, it must be updated", {"name": name}, ) continue logging.info("Get token with %s remaining requests", rest) if best_token is None: best_token = Token(user, value, rest) elif best_token.rest < rest: best_token = Token(user, value, rest) if best_token.rest > SAFE_REQUESTS_LIMIT: break assert best_token ROBOT_TOKEN = best_token logging.info( "User %s with %s remaining requests is used", ROBOT_TOKEN.user.login, ROBOT_TOKEN.rest, ) return ROBOT_TOKEN.value