#!/usr/bin/python

# file:        google_api_script.py
# version:     1.0
# author:      ben servoz
# description: groundwork for a class handling Google API interaction
# notes:       instructions at http://ben.akrin.com/?p=3923

import base64
import json
import math
import os.path
import pprint
import pycurl
import random
import signal
import socket
from StringIO import StringIO
import sys
from sys import argv
import time
import urllib

# boolean used to terminate threads if SIGINT was received
sigint_received = False ;
def signal_handler( signal, frame ):
	global sigint_received
	sys.stdout.write( "SIGINT received, cleaning up threads and bailing." ) ;
	sigint_received = True ;
signal.signal(signal.SIGINT, signal_handler)


class google_api:
	"class to handle Google Drive interaction"
	access_tokens = {}


	def __init__( self, domain_account_json_file ):
		# loading JSON file config
		domain_account_config_handle = open( domain_account_json_file )
		domain_account_config = json.load( domain_account_config_handle )
		self.client_id = domain_account_config['client_id']
		self.client_email = domain_account_config['client_email']
		self.private_key = domain_account_config['private_key']


	def get_about( self, sub_as ):
		request_url = "https://www.googleapis.com/drive/v2/about"
		request_type = "GET"
		request_get = ""
		request_post = ""
		request_headers = []
		request_body = ""

		result = self.make_api_call( request_url, request_type, request_get, request_post, request_headers, request_body, sub_as )
		if result['response_code']!=200:
			if result['response_code']==404:
				sys.stderr.write( "WARNING: it looks like the resource went away before I had a chance to get_about\n" )
				return False
			else:
				sys.stderr.write( "ERROR: didn't get a 200 back from Google trying to get_about\n" )
				global fatal_error_reached # if I don't put it everywhere I have an error I'll forget and it won't kick in, nice job Python.
				fatal_error_reached = True ;
		try:
			data = json.loads( result['response_body'] )
		except:
			print "THIS IS THE DATA" + result['response_body'] + "|||"
			print "THIS IS THE DATA" + str(len(result['response_body'])) + "|||"
			print "THIS IS THE DATA" + type(result['response_body'])+ "|||"
		return data


	def refresh_access_token( self, sub_as=None ):
		# building JWT
		jwt_request_date = time.time()
		jwt_expiration_date = jwt_request_date + 60*60 # +1 hour
		jwt_header = '{"alg":"RS256","typ":"JWT"}' ;
		jwt_claim_set = '{"iss":"' + self.client_email + '","scope":"https://www.googleapis.com/auth/drive","aud":"https://www.googleapis.com/oauth2/v3/token","exp":' + str(jwt_expiration_date) + ',"iat":' + str(jwt_request_date) ;
		if sub_as!=None and sub_as!=self.client_email:
			jwt_claim_set += ',"sub":"' + sub_as + '"'
		jwt_claim_set += '}'

		jwt_signature = self.sign_data( self.private_key, base64.urlsafe_b64encode(jwt_header) + '.' + base64.urlsafe_b64encode(jwt_claim_set) )
		jwt = base64.urlsafe_b64encode( jwt_header ) + "." + base64.urlsafe_b64encode( jwt_claim_set ) + "." + base64.urlsafe_b64encode( jwt_signature ) 

		request_url = "https://www.googleapis.com/oauth2/v3/token"
		request_type = "POST"
		request_get = ""
		request_post = "grant_type=" + urllib.quote_plus( "urn:ietf:params:oauth:grant-type:jwt-bearer" ) + "&assertion=" + jwt
		request_headers = ["application/x-www-form-urlencoded"]
		request_body = ""

		result = self.make_api_call( request_url, request_type, request_get, request_post, request_headers, request_body )
		if result['response_code']!=200:
			sys.stderr.write( "ERROR: didn't get a 200 back from Google trying to refresh_access_token\n" )
			global fatal_error_reached # if I don't put it everywhere I have an error I'll forget and it won't kick in, nice job Python.
			fatal_error_reached = True ;
		access_token = json.loads( result['response_body'] )
		if sub_as==None:
			sub_as = self.client_email
		self.access_tokens[sub_as] = {'token':access_token, 'expires_at':jwt_expiration_date}


	def sign_data( self, private_key, data ):
	    from Crypto.PublicKey import RSA
	    from Crypto.Signature import PKCS1_v1_5
	    from Crypto.Hash import SHA256
	    from base64 import b64encode, b64decode
	    rsakey = RSA.importKey( private_key )
	    signer = PKCS1_v1_5.new( rsakey )
	    digest = SHA256.new()
	    digest.update( data )
	    sign = signer.sign( digest )
	    return sign


	def make_api_call( self, request_url, request_type, request_get, request_post, request_headers, request_body, inject_token_of=None, exponential_backoff_counter=None ):
		
		# getting an access token if we don't have one already
		if inject_token_of!=None and inject_token_of not in self.access_tokens.keys():
			self.refresh_access_token( inject_token_of )
		# let's try and be proactive about refreshing tokens right before they expire
		if inject_token_of!=None and inject_token_of in self.access_tokens.keys() and (time.time()+10)>self.access_tokens[inject_token_of]['expires_at']:
			self.refresh_access_token( inject_token_of )

		body = StringIO()
		headers = StringIO()
		c = pycurl.Curl()
		c.setopt( c.CONNECTTIMEOUT, 5 )
		c.setopt( c.TIMEOUT, 5 )
		c.setopt( c.NOSIGNAL, 5 )
		if request_get==None or request_get=="":
			c.setopt( c.URL, request_url.encode("ascii") )
		else:
			c.setopt( c.URL, request_url.encode("ascii") + "?" + request_get.encode("ascii") )
		if not( request_post==None or request_post=="" ):
			c.setopt( c.POSTFIELDS, request_post.encode("ascii") )
			c.setopt( c.POST, 1 )
		c.setopt( c.FOLLOWLOCATION, True )
		c.setopt( c.CUSTOMREQUEST, request_type )
		add_to_headers = ["Content-length: " +  str(len(request_post))]
		if inject_token_of!=None:
			add_to_headers.append( "Authorization: Bearer " + self.access_tokens[inject_token_of]['token']['access_token'].encode("ascii") )
		c.setopt( c.HTTPHEADER, request_headers + add_to_headers )
		c.setopt( c.WRITEFUNCTION, body.write )
		c.setopt( c.HEADERFUNCTION, headers.write )
		try:
			c.perform()
		except: 
			pass
		response_code = c.getinfo( c.RESPONSE_CODE )
		response_body = body.getvalue()
		response_headers = headers.getvalue()
		c.close()
		body.close()
		headers.close()

		if response_code==200 or response_code==204 or response_code==404: # 200 or 204 when everything is cool; 404 when the resource went away before we had a chance to do anything with it. We pass that back so the originating functions get a chances to handle it properly.
			return {'response_code':response_code, 'response_body':response_body, 'response_headers':response_headers}
		elif response_code==401:
			if inject_token_of==None:
				sys.stderr.write( "ERROR: I've gotten a 401 back from Google when I wasn't even supposed to be injecting a token; something's gone horribly wrong.\n" ) ;
				sys.exit( 1 )
			else:
				sys.stderr.write( "WARNING: the logic should be set so that we shouldn't get 401s back from Google\n" )
				self.refresh_access_token( inject_token_of )
				# and then we try again
				return self.make_api_call( request_url, request_type, request_get, request_post, request_headers, request_body, inject_token_of, exponential_backoff_counter )
		else:
			sys.stdout.write( "exponential backoff with http response code " + str(response_code) + " url:" + request_url + " body:" + response_body + "body size:" + str(len(response_body)) + "\n" )
			# in any other case, we implement exponential backoff (https://developers.google.com/drive/web/handle-errors)
			if exponential_backoff_counter==None:
				exponential_backoff_counter = 0
			exponential_backoff_counter += 1
			# Google say to exponential back-off up to 5 but I've seen enough weird things that we do up to ~2 hours
			if exponential_backoff_counter>12:
				sys.stdout.write( "exponential backoff limit exceeded, gotta return...\n" )
				# well, we've tried and tried to no avail, time to return what we are getting
				return {'response_code':response_code, 'response_body':response_body, 'response_headers':response_headers}
			else:
				sleep_for = math.pow( 2, exponential_backoff_counter )
				while sleep_for>0:
					if sigint_received:
						return False ;
					time.sleep( 1 )
					sleep_for = sleep_for - 1
				return self.make_api_call( request_url, request_type, request_get, request_post, request_headers, request_body, inject_token_of, exponential_backoff_counter )


script, domain_account_json_file, account_to_sub_as = argv
gapi = google_api( domain_account_json_file )
result = gapi.get_about( account_to_sub_as )
pp = pprint.PrettyPrinter( indent=4 )
pp.pprint( result )
sys.exit( 0 )
