Spaces:

Text-to-Document-Generation
/

Docgenie-API

Paused

Docgenie-API / api /test_get_google_token.py

Ahadhassan-2003

deploy: update HF Space

dc4e6da 6 months ago

8.51 kB

"""
Helper script to get Google Drive OAuth token for testing.
This script implements the OAuth flow to get access and refresh tokens
from Google Drive API for testing purposes.
Prerequisites:
1. Google Cloud Project with Drive API enabled
2. OAuth 2.0 Client ID credentials
3. Add http://localhost:8080 as authorized redirect URI
Usage:
    python test_get_google_token.py --client-id YOUR_CLIENT_ID --client-secret YOUR_CLIENT_SECRET
"""

import argparse
import webbrowser
from urllib.parse import urlencode, parse_qs
from http.server import HTTPServer, BaseHTTPRequestHandler
import requests


# Global variable to store authorization code
auth_code = None


class OAuthCallbackHandler(BaseHTTPRequestHandler):
    """HTTP server to handle OAuth callback"""
    
    def do_GET(self):
        global auth_code
        
        # Parse query parameters
        query = self.path.split('?', 1)[-1]
        params = parse_qs(query)
        
        if 'code' in params:
            auth_code = params['code'][0]
            
            # Send success response
            self.send_response(200)
            self.send_header('Content-type', 'text/html')
            self.end_headers()
            
            html = """
            <html>
            <head><title>Authorization Successful</title></head>
            <body style="font-family: Arial; text-align: center; padding: 50px;">
                <h1 style="color: green;">✓ Authorization Successful!</h1>
                <p>You can close this window and return to the terminal.</p>
            </body>
            </html>
            """
            self.wfile.write(html.encode())
        else:
            # Error response
            self.send_response(400)
            self.send_header('Content-type', 'text/html')
            self.end_headers()
            
            error = params.get('error', ['Unknown error'])[0]
            html = f"""
            <html>
            <head><title>Authorization Failed</title></head>
            <body style="font-family: Arial; text-align: center; padding: 50px;">
                <h1 style="color: red;">✗ Authorization Failed</h1>
                <p>Error: {error}</p>
                <p>Please try again.</p>
            </body>
            </html>
            """
            self.wfile.write(html.encode())
    
    def log_message(self, format, *args):
        """Suppress default logging"""
        pass


def get_google_drive_token(client_id: str, client_secret: str, redirect_uri: str = "http://localhost:8080"):
    """
    Get Google Drive OAuth tokens through OAuth flow.
    
    Args:
        client_id: Google OAuth client ID
        client_secret: Google OAuth client secret
        redirect_uri: OAuth redirect URI (must match Google Cloud Console)
    
    Returns:
        dict with 'access_token' and 'refresh_token'
    """
    global auth_code
    
    print("=" * 80)
    print(" " * 20 + "GOOGLE DRIVE OAUTH TOKEN GENERATOR")
    print("=" * 80)
    print()
    
    # Step 1: Generate authorization URL
    auth_params = {
        'client_id': client_id,
        'redirect_uri': redirect_uri,
        'response_type': 'code',
        'scope': 'https://www.googleapis.com/auth/drive.file',
        'access_type': 'offline',  # Get refresh token
        'prompt': 'consent'  # Force consent to get refresh token
    }
    
    auth_url = f"https://accounts.google.com/o/oauth2/v2/auth?{urlencode(auth_params)}"
    
    print("Step 1: Authorize with Google")
    print("-" * 80)
    print("\nOpening authorization URL in your browser...")
    print("If it doesn't open automatically, copy this URL:\n")
    print(auth_url)
    print()
    
    # Open browser
    webbrowser.open(auth_url)
    
    # Step 2: Start local server to receive callback
    print("Step 2: Waiting for authorization...")
    print("-" * 80)
    print(f"Local server listening on {redirect_uri}")
    print("Complete the authorization in your browser.")
    print()
    
    server = HTTPServer(('localhost', 8080), OAuthCallbackHandler)
    
    # Wait for one request (the callback)
    while auth_code is None:
        server.handle_request()
    
    server.server_close()
    
    if not auth_code:
        print("✗ Failed to get authorization code")
        return None
    
    print("✓ Authorization code received!")
    print()
    
    # Step 3: Exchange code for tokens
    print("Step 3: Exchanging code for tokens...")
    print("-" * 80)
    
    token_url = "https://oauth2.googleapis.com/token"
    token_data = {
        'code': auth_code,
        'client_id': client_id,
        'client_secret': client_secret,
        'redirect_uri': redirect_uri,
        'grant_type': 'authorization_code'
    }
    
    try:
        response = requests.post(token_url, data=token_data)
        response.raise_for_status()
        tokens = response.json()
        
        print("✓ Tokens received!")
        print()
        print("=" * 80)
        print(" " * 30 + "TOKENS")
        print("=" * 80)
        print()
        print("Access Token:")
        print(tokens['access_token'])
        print()
        
        if 'refresh_token' in tokens:
            print("Refresh Token:")
            print(tokens['refresh_token'])
            print()
        else:
            print("⚠ No refresh token received (user may have authorized before)")
            print("  To get a refresh token:")
            print("  1. Go to: https://myaccount.google.com/permissions")
            print("  2. Remove your app's access")
            print("  3. Run this script again")
            print()
        
        print("Expires In: {} seconds".format(tokens.get('expires_in', 'N/A')))
        print()
        
        # Show usage instructions
        print("=" * 80)
        print(" " * 25 + "USAGE INSTRUCTIONS")
        print("=" * 80)
        print()
        print("Option 1: Use with test script directly")
        print("-" * 80)
        print("python test_async_api.py \\")
        print(f"  --google-token {tokens['access_token']}")
        if 'refresh_token' in tokens:
            print(f"  --google-refresh-token {tokens['refresh_token']}")
        print()
        
        print("Option 2: Set environment variable")
        print("-" * 80)
        print(f"export GOOGLE_DRIVE_TOKEN=\"{tokens['access_token']}\"")
        if 'refresh_token' in tokens:
            print(f"export GOOGLE_DRIVE_REFRESH_TOKEN=\"{tokens['refresh_token']}\"")
        print("python test_async_api.py")
        print()
        
        print("Option 3: Use in your frontend")
        print("-" * 80)
        print("Store these tokens in your frontend application and include them")
        print("in API requests to /generate/async endpoint.")
        print()
        
        print("=" * 80)
        
        return tokens
        
    except Exception as e:
        print(f"✗ Failed to exchange code for tokens: {e}")
        if hasattr(e, 'response') and e.response:
            print(f"Response: {e.response.text}")
        return None


def main():
    parser = argparse.ArgumentParser(
        description="Get Google Drive OAuth token for testing"
    )
    parser.add_argument(
        "--client-id",
        type=str,
        required=True,
        help="Google OAuth Client ID"
    )
    parser.add_argument(
        "--client-secret",
        type=str,
        required=True,
        help="Google OAuth Client Secret"
    )
    parser.add_argument(
        "--redirect-uri",
        type=str,
        default="http://localhost:8080",
        help="OAuth redirect URI (default: http://localhost:8080)"
    )
    
    args = parser.parse_args()
    
    print()
    print("Prerequisites Check:")
    print("-" * 80)
    print(f"✓ Client ID: {args.client_id[:20]}...")
    print(f"✓ Client Secret: {args.client_secret[:10]}...")
    print(f"✓ Redirect URI: {args.redirect_uri}")
    print()
    print("Make sure you've added this redirect URI to your Google Cloud Console:")
    print("  https://console.cloud.google.com/apis/credentials")
    print()
    input("Press Enter to continue...")
    print()
    
    tokens = get_google_drive_token(
        client_id=args.client_id,
        client_secret=args.client_secret,
        redirect_uri=args.redirect_uri
    )
    
    if tokens:
        print("✓ SUCCESS! Use the tokens above to test the async API.")
    else:
        print("✗ FAILED to get tokens. Please check your credentials and try again.")


if __name__ == "__main__":
    main()

	"""
	Helper script to get Google Drive OAuth token for testing.

	This script implements the OAuth flow to get access and refresh tokens
	from Google Drive API for testing purposes.

	Prerequisites:
	1. Google Cloud Project with Drive API enabled
	2. OAuth 2.0 Client ID credentials
	3. Add http://localhost:8080 as authorized redirect URI

	Usage:
	python test_get_google_token.py --client-id YOUR_CLIENT_ID --client-secret YOUR_CLIENT_SECRET
	"""

	import argparse
	import webbrowser
	from urllib.parse import urlencode, parse_qs
	from http.server import HTTPServer, BaseHTTPRequestHandler
	import requests


	# Global variable to store authorization code
	auth_code = None


	class OAuthCallbackHandler(BaseHTTPRequestHandler):
	"""HTTP server to handle OAuth callback"""

	def do_GET(self):
	global auth_code

	# Parse query parameters
	query = self.path.split('?', 1)[-1]
	params = parse_qs(query)

	if 'code' in params:
	auth_code = params['code'][0]

	# Send success response
	self.send_response(200)
	self.send_header('Content-type', 'text/html')
	self.end_headers()

	html = """
	<html>
	<head><title>Authorization Successful</title></head>
	<body style="font-family: Arial; text-align: center; padding: 50px;">
	<h1 style="color: green;">✓ Authorization Successful!</h1>
	<p>You can close this window and return to the terminal.</p>
	</body>
	</html>
	"""
	self.wfile.write(html.encode())
	else:
	# Error response
	self.send_response(400)
	self.send_header('Content-type', 'text/html')
	self.end_headers()

	error = params.get('error', ['Unknown error'])[0]
	html = f"""
	<html>
	<head><title>Authorization Failed</title></head>
	<body style="font-family: Arial; text-align: center; padding: 50px;">
	<h1 style="color: red;">✗ Authorization Failed</h1>
	<p>Error: {error}</p>
	<p>Please try again.</p>
	</body>
	</html>
	"""
	self.wfile.write(html.encode())

	def log_message(self, format, *args):
	"""Suppress default logging"""
	pass


	def get_google_drive_token(client_id: str, client_secret: str, redirect_uri: str = "http://localhost:8080"):
	"""
	Get Google Drive OAuth tokens through OAuth flow.

	Args:
	client_id: Google OAuth client ID
	client_secret: Google OAuth client secret
	redirect_uri: OAuth redirect URI (must match Google Cloud Console)

	Returns:
	dict with 'access_token' and 'refresh_token'
	"""
	global auth_code

	print("=" * 80)
	print(" " * 20 + "GOOGLE DRIVE OAUTH TOKEN GENERATOR")
	print("=" * 80)
	print()

	# Step 1: Generate authorization URL
	auth_params = {
	'client_id': client_id,
	'redirect_uri': redirect_uri,
	'response_type': 'code',
	'scope': 'https://www.googleapis.com/auth/drive.file',
	'access_type': 'offline', # Get refresh token
	'prompt': 'consent' # Force consent to get refresh token
	}

	auth_url = f"https://accounts.google.com/o/oauth2/v2/auth?{urlencode(auth_params)}"

	print("Step 1: Authorize with Google")
	print("-" * 80)
	print("\nOpening authorization URL in your browser...")
	print("If it doesn't open automatically, copy this URL:\n")
	print(auth_url)
	print()

	# Open browser
	webbrowser.open(auth_url)

	# Step 2: Start local server to receive callback
	print("Step 2: Waiting for authorization...")
	print("-" * 80)
	print(f"Local server listening on {redirect_uri}")
	print("Complete the authorization in your browser.")
	print()

	server = HTTPServer(('localhost', 8080), OAuthCallbackHandler)

	# Wait for one request (the callback)
	while auth_code is None:
	server.handle_request()

	server.server_close()

	if not auth_code:
	print("✗ Failed to get authorization code")
	return None

	print("✓ Authorization code received!")
	print()

	# Step 3: Exchange code for tokens
	print("Step 3: Exchanging code for tokens...")
	print("-" * 80)

	token_url = "https://oauth2.googleapis.com/token"
	token_data = {
	'code': auth_code,
	'client_id': client_id,
	'client_secret': client_secret,
	'redirect_uri': redirect_uri,
	'grant_type': 'authorization_code'
	}

	try:
	response = requests.post(token_url, data=token_data)
	response.raise_for_status()
	tokens = response.json()

	print("✓ Tokens received!")
	print()
	print("=" * 80)
	print(" " * 30 + "TOKENS")
	print("=" * 80)
	print()
	print("Access Token:")
	print(tokens['access_token'])
	print()

	if 'refresh_token' in tokens:
	print("Refresh Token:")
	print(tokens['refresh_token'])
	print()
	else:
	print("⚠ No refresh token received (user may have authorized before)")
	print(" To get a refresh token:")
	print(" 1. Go to: https://myaccount.google.com/permissions")
	print(" 2. Remove your app's access")
	print(" 3. Run this script again")
	print()

	print("Expires In: {} seconds".format(tokens.get('expires_in', 'N/A')))
	print()

	# Show usage instructions
	print("=" * 80)
	print(" " * 25 + "USAGE INSTRUCTIONS")
	print("=" * 80)
	print()
	print("Option 1: Use with test script directly")
	print("-" * 80)
	print("python test_async_api.py \\")
	print(f" --google-token {tokens['access_token']}")
	if 'refresh_token' in tokens:
	print(f" --google-refresh-token {tokens['refresh_token']}")
	print()

	print("Option 2: Set environment variable")
	print("-" * 80)
	print(f"export GOOGLE_DRIVE_TOKEN=\"{tokens['access_token']}\"")
	if 'refresh_token' in tokens:
	print(f"export GOOGLE_DRIVE_REFRESH_TOKEN=\"{tokens['refresh_token']}\"")
	print("python test_async_api.py")
	print()

	print("Option 3: Use in your frontend")
	print("-" * 80)
	print("Store these tokens in your frontend application and include them")
	print("in API requests to /generate/async endpoint.")
	print()

	print("=" * 80)

	return tokens

	except Exception as e:
	print(f"✗ Failed to exchange code for tokens: {e}")
	if hasattr(e, 'response') and e.response:
	print(f"Response: {e.response.text}")
	return None


	def main():
	parser = argparse.ArgumentParser(
	description="Get Google Drive OAuth token for testing"
	)
	parser.add_argument(
	"--client-id",
	type=str,
	required=True,
	help="Google OAuth Client ID"
	)
	parser.add_argument(
	"--client-secret",
	type=str,
	required=True,
	help="Google OAuth Client Secret"
	)
	parser.add_argument(
	"--redirect-uri",
	type=str,
	default="http://localhost:8080",
	help="OAuth redirect URI (default: http://localhost:8080)"
	)

	args = parser.parse_args()

	print()
	print("Prerequisites Check:")
	print("-" * 80)
	print(f"✓ Client ID: {args.client_id[:20]}...")
	print(f"✓ Client Secret: {args.client_secret[:10]}...")
	print(f"✓ Redirect URI: {args.redirect_uri}")
	print()
	print("Make sure you've added this redirect URI to your Google Cloud Console:")
	print(" https://console.cloud.google.com/apis/credentials")
	print()
	input("Press Enter to continue...")
	print()

	tokens = get_google_drive_token(
	client_id=args.client_id,
	client_secret=args.client_secret,
	redirect_uri=args.redirect_uri
	)

	if tokens:
	print("✓ SUCCESS! Use the tokens above to test the async API.")
	else:
	print("✗ FAILED to get tokens. Please check your credentials and try again.")


	if __name__ == "__main__":
	main()