Skip to content

Commit

Permalink
authenticate CLI with agentstack account
Browse files Browse the repository at this point in the history
  • Loading branch information
bboynton97 committed Dec 22, 2024
1 parent 69bd4d5 commit 8cef7a8
Show file tree
Hide file tree
Showing 2 changed files with 148 additions and 1 deletion.
142 changes: 142 additions & 0 deletions agentstack/auth.py
Original file line number Diff line number Diff line change
@@ -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 = """
<html>
<body>
<script>
setTimeout(function() {
window.close();
}, 1000);
</script>
<h2>Authentication successful! You can close this window.</h2>
</body>
</html>
"""
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
7 changes: 6 additions & 1 deletion agentstack/main.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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]
Expand Down Expand Up @@ -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:
Expand Down

0 comments on commit 8cef7a8

Please sign in to comment.