<reponame>rohit-k-das/vulnerability-management-reporter<gh_stars>0
import requests
from requests.adapters import HTTPAdapter
from requests.packages.urllib3.util.retry import Retry
import json
import logging
import time
import jwt
from typing import List, Dict, Tuple
import ConfigParser
import os

MAX_THREADS = 14  # Get max number of threads for multi-threading

logger = logging.getLogger(__name__)

Config = ConfigParser.ConfigParser()
Config.read(os.path.join(os.path.abspath(os.path.dirname(__file__)), 'settings.ini'))
google_service_account_secret = Config.get('Settings', 'Google_Service_Account_Secret')
google_service_account_id = Config.get('Settings', 'Google_Service_Account_ID')

sheet_api = 'https://sheets.googleapis.com/v4/{0}'
sheet_access_tokens = {}


# Generate session with max of 3 retries and interval of 1 second
def session_generator():
    session = requests.Session()
    retry = Retry(connect=3, backoff_factor=0.5)
    adapter = HTTPAdapter(max_retries=retry)
    session.mount('http://', adapter)
    session.mount('https://', adapter)
    return session


# Create OAuth token per requirement for each recipient
def generate_sheet_api_access_token(recipient: str) -> Tuple[str, int]:
    access_token = None
    expiry = None
    jwt_header = {"alg": "RS256", "typ": "JWT"}
    iat = time.time()
    exp = iat + 3600
    jwt_claim_set = {
        'iss': google_service_account_id,
        'scope': 'https://www.googleapis.com/auth/drive.file https://www.googleapis.com/auth/spreadsheets',
        'sub': recipient,
        'aud': 'https://www.googleapis.com/oauth2/v4/token',
        'iat': iat,
        'exp': exp
    }

    secret = bytes(google_service_account_secret.replace('\\n', '\n'), 'utf-8')
    signed_jwt = jwt.encode(jwt_claim_set, secret, headers=jwt_header, algorithm='RS256')

    headers = {"Content-Type": "application/json; charset=utf-8"}
    data = {'grant_type': 'urn:ietf:params:oauth:grant-type:jwt-bearer', 'assertion': signed_jwt.decode('utf-8').replace("'", '"')}
    url = 'https://www.googleapis.com/oauth2/v4/token'
    session = session_generator()
    resp = session.post(url, headers=headers, data=json.dumps(data))
    response = resp.json()
    if resp.ok:
        access_token = response['access_token']
        expiry = time.time() + response['expires_in']
    elif resp.status_code == 429:
        logger.error('Too many requests. Sleeping %s' % response['error_description'])
        time.sleep(1)
        access_token, expiry = generate_sheet_api_access_token(recipient)
    else:
        logger.error('Failed to generate access token')
        logger.error("%d:%s" % (resp.status_code, resp.text))
    return access_token, expiry


def create_sheet(filename: str, sheet_content: Dict[str, List[str]], recipient: str, sheet_id: str=None) -> str:
    id = None
    access_token = sheet_access_tokens[recipient]['access_token']
    expiry = sheet_access_tokens[recipient]['expiry']
    query_start_time = time.time()

    sheets = []
    for each_sheet in sheet_content:
        sheet_dict = {'properties': {'title': each_sheet}, 'data': []}
        rows_dict = {'rowData': []}
        for each_row in sheet_content[each_sheet]:
            values = []
            for each_item in each_row:
                value_dict = {}
                value_dict['userEnteredValue'] = {}
                value_dict['userEnteredValue']['stringValue'] = each_item
                values.append(value_dict)
            rows_dict['rowData'].append({'values': values})
        sheet_dict['data'].append(rows_dict)
        sheets.append(sheet_dict)

    # Make the API call if token expiry time is greater than 1 minute
    if (expiry - query_start_time) > 60:
        headers = {'Content-Type': 'application/json', 'Authorization': 'Bearer %s' % access_token}
        data = {
            'sheets': sheets
        }

        if sheet_id is not None:
            data['spreadsheetId'] = sheet_id

        url = sheet_api.format('spreadsheets')
        session = session_generator()
        resp = session.post(url, headers=headers, json=data)
        if resp.ok:
            response = resp.json()
            if 'spreadsheetId' in response:
                id = response['spreadsheetId']
        elif resp.status_code == 429:
            logger.error('Too many requests. Sleeping %s' % resp.json()['error']['message'])
            time.sleep(1)
            id = create_sheet(filename, sheet_content, recipient, sheet_id)
        else:
            logger.error('Failed to create google sheet %s' % filename)
            logger.error("%d:%s" % (resp.status_code, resp.text))
    # Create new access token to be used by the recipient
    else:
        access_token, expiry = generate_sheet_api_access_token(recipient)
        if access_token is not None and expiry is not None:
            sheet_access_tokens[recipient]['access_token'] = access_token
            sheet_access_tokens[recipient]['expiry'] = expiry
            id = create_sheet(filename, sheet_content, recipient, sheet_id)
    return id


def update_sheet(filename: str, sheet_content: List[List[str]], recipient: str, sheet_id: str) -> None:
    access_token = sheet_access_tokens[recipient]['access_token']
    expiry = sheet_access_tokens[recipient]['expiry']
    query_start_time = time.time()
    params = {'valueInputOption': 'USER_ENTERED'}

    # Make the API call if token expiry time is greater than 1 minute
    if (expiry - query_start_time) > 60:
        headers = {'Content-Type': 'application/json', 'Authorization': 'Bearer %s' % access_token}
        data = {
            'values': sheet_content
        }

        url = sheet_api.format('spreadsheets/%s/values/A1' % sheet_id)
        session = session_generator()
        resp = session.put(url, headers=headers, json=data, params=params)
        response = resp.json()
        if resp.ok:
            if response['updatedRows'] > 0:
                logger.info('Updated sheet %s' % filename)
        elif resp.status_code == 429:
            logger.error('Too many requests. Sleeping %s' % response['error']['message'])
            time.sleep(1)
            update_sheet(filename, sheet_content, recipient, sheet_id)
        elif resp.status_code == 400 and 'Your input contains more than the maximum of 50000 characters in a single cell' in response['error']['message']:
            logger.warning('Data limit(50,000) exceeded in a single cell. Dumping to file %s' % filename)
            with open(filename, 'w') as f:
                f.write(str(sheet_content))
        elif resp.status_code == 500:
            time.sleep(30)
            logger.warning('Encountered 500. Retrying %s' % filename)
            update_sheet(filename, sheet_content, recipient, sheet_id)
        else:
            logger.error('Failed to update google sheet %s' % filename)
            logger.error(json.dumps(response, indent=2))
    # Create new access token to be used by the recipient
    else:
        access_token, expiry = generate_sheet_api_access_token(recipient)
        if access_token is not None and expiry is not None:
            sheet_access_tokens[recipient]['access_token'] = access_token
            sheet_access_tokens[recipient]['expiry'] = expiry
            update_sheet(filename, sheet_content, recipient, sheet_id)


def append_sheet(filename: str, sheet_content: List[List[str]], recipient: str, sheet_id: str) -> None:
    access_token = sheet_access_tokens[recipient]['access_token']
    expiry = sheet_access_tokens[recipient]['expiry']
    query_start_time = time.time()
    params = {'valueInputOption': 'USER_ENTERED', 'insertDataOption': 'INSERT_ROWS'}

    # Make the API call if token expiry time is greater than 1 minute
    if (expiry - query_start_time) > 60:
        headers = {'Content-Type': 'application/json', 'Authorization': 'Bearer %s' % access_token}
        data = {
            'values': sheet_content
        }

        url = sheet_api.format('spreadsheets/%s/values/A1:append' % sheet_id)
        session = session_generator()
        resp = session.post(url, headers=headers, json=data, params=params)
        if resp.ok:
            response = resp.json()
            if response['updates']['updatedRows'] > 0:
                logger.info('Updated sheet %s' % filename)
        elif resp.status_code == 429:
            logger.error('Too many requests. Sleeping %s' % resp.json()['error']['message'])
            time.sleep(1)
            append_sheet(filename, sheet_content, recipient, sheet_id)
        elif resp.status_code == 400 and 'Your input contains more than the maximum of 50000 characters in a single cell' in resp.json()['error']['message']:
            logger.warning('Data limit(50,000) exceeded in a single cell. Dumping to file %s' % filename)
            with open(filename, 'w') as f:
                f.write(str(sheet_content))
        else:
            logger.error('Failed to append google sheet %s' % filename)
            logger.error("%d:%s" % (resp.status_code, resp.text))
    # Create new access token to be used by the recipient
    else:
        access_token, expiry = generate_sheet_api_access_token(recipient)
        if access_token is not None and expiry is not None:
            sheet_access_tokens[recipient]['access_token'] = access_token
            sheet_access_tokens[recipient]['expiry'] = expiry
            append_sheet(filename, sheet_content, recipient, sheet_id)
