From 8cef7a882c67384ea4724f922bdec38767f91875 Mon Sep 17 00:00:00 2001 From: Braelyn Boynton Date: Sun, 22 Dec 2024 15:48:37 -0500 Subject: [PATCH] authenticate CLI with agentstack account --- agentstack/auth.py | 142 +++++++++++++++++++++++++++++++++++++++++++++ agentstack/main.py | 7 ++- 2 files changed, 148 insertions(+), 1 deletion(-) create mode 100644 agentstack/auth.py diff --git a/agentstack/auth.py b/agentstack/auth.py new file mode 100644 index 0000000..a4fc46c --- /dev/null +++ b/agentstack/auth.py @@ -0,0 +1,142 @@ +from http.server import HTTPServer, BaseHTTPRequestHandler +from urllib.parse import parse_qs, urlparse +import webbrowser +import json +import os +import threading +import socket +from pathlib import Path + +import inquirer +from appdirs import user_data_dir +from agentstack.logger import log + + +try: + base_dir = Path(user_data_dir("agentstack", "agency")) + # Test if we can write to directory + test_file = base_dir / '.test_write_permission' + test_file.touch() + test_file.unlink() +except (RuntimeError, OSError, PermissionError): + # In CI or when directory is not writable, use temp directory + base_dir = Path(os.getenv('TEMP', '/tmp')) + + +class AuthCallbackHandler(BaseHTTPRequestHandler): + def do_GET(self): + """Handle the OAuth callback from the browser""" + try: + # Parse the query parameters + query_components = parse_qs(urlparse(self.path).query) + + # Extract the token from query parameters + token = query_components.get('token', [''])[0] + + if token: + # Store the token + base_dir.mkdir(exist_ok=True, parents=True) + + with open(base_dir / 'auth.json', 'w') as f: + json.dump({'bearer_token': token}, f) + + # Send success response + self.send_response(200) + self.send_header('Content-type', 'text/html') + self.end_headers() + + success_html = """ + + + +

Authentication successful! You can close this window.

+ + + """ + self.wfile.write(success_html.encode()) + + # Signal the main thread that we're done + self.server.authentication_successful = True + else: + self.send_response(400) + self.send_header('Content-type', 'text/html') + self.end_headers() + self.wfile.write(b'Authentication failed: No token received') + + except Exception as e: + self.send_response(500) + self.send_header('Content-type', 'text/html') + self.end_headers() + self.wfile.write(f'Error: {str(e)}'.encode()) + +def find_free_port(): + """Find a free port on localhost""" + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind(('', 0)) + s.listen(1) + port = s.getsockname()[1] + return port + +def start_auth_server(): + """Start the local authentication server""" + port = find_free_port() + server = HTTPServer(('localhost', port), AuthCallbackHandler) + server.authentication_successful = False + return server, port + + +def login(): + """Log in to AgentStack""" + try: + # check if already logged in + token = get_stored_token() + if token: + print("You are already authenticated!") + if not inquirer.confirm('Would you like to log in with a different account?'): + return + + # Start the local server + server, port = start_auth_server() + + # Create server thread + server_thread = threading.Thread(target=server.serve_forever) + server_thread.daemon = True + server_thread.start() + + # Open the browser to the login page + auth_url_base = os.getenv('AGENTSTACK_AUTHORIZATION_BASE_URL', 'https://agentstack.sh') + auth_url = f"{auth_url_base}/login?callback_port={port}" + webbrowser.open(auth_url) + + # Wait for authentication to complete + while not server.authentication_successful: + pass + + # Cleanup + server.shutdown() + server_thread.join() + + print("🔐 Authentication successful! Token has been stored.") + return True + + except Exception as e: + log.warn(f"Authentication failed: {str(e)}", err=True) + return False + + +def get_stored_token(): + """Retrieve the stored bearer token""" + try: + auth_path = base_dir / 'auth.json' + if not auth_path.exists(): + return None + + with open(auth_path) as f: + config = json.load(f) + return config.get('bearer_token') + except Exception: + return None \ No newline at end of file diff --git a/agentstack/main.py b/agentstack/main.py index 70be8da..07e0c97 100644 --- a/agentstack/main.py +++ b/agentstack/main.py @@ -2,7 +2,7 @@ import argparse import webbrowser -from agentstack import conf +from agentstack import conf, auth from agentstack.cli import ( init_project_builder, add_tool, @@ -50,6 +50,9 @@ def main(): # 'templates' command subparsers.add_parser("templates", help="View Agentstack templates") + # 'login' command + subparsers.add_parser("login", help="Authenticate with Agentstack.sh") + # 'init' command init_parser = subparsers.add_parser( "init", aliases=["i"], help="Initialize a directory for the project", parents=[global_parser] @@ -188,6 +191,8 @@ def main(): tools_parser.print_help() elif args.command in ['export', 'e']: export_template(args.filename) + elif args.command in ['login']: + auth.login() elif args.command in ['update', 'u']: pass # Update check already done else: