AgentStack/agentstack/auth.py at d388e06c9190667b8a436f4af39674f491f3e652 · agentstack-ai/AgentStack · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
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 questionary
from appdirs import user_data_dir
from agentstack 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:
            log.success("You are already authenticated!")
            if not questionary.confirm('Would you like to log in with a different account?').ask():
                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()

        log.success("🔐 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