Feature: simple login to prime (#35)

* simple login to prime using `prime login`

Author: JannikSt

pyproject.toml

@@ -12,7 +12,8 @@ dependencies = [
     "typer[all]>=0.9.0",
     "requests>=2.31.0",
     "rich>=13.3.1",  
-    "pydantic>=2.0.0" 
+    "pydantic>=2.0.0",
+    "cryptography>=41.0.0"
 ]
 keywords = ["cli", "gpu", "cloud", "compute"]
 classifiers = [

src/prime_cli/commands/login.py

@@ -0,0 +1,145 @@
+import base64
+import time
+from typing import Optional
+
+import requests
+import typer
+from cryptography.hazmat.primitives import hashes, serialization
+from cryptography.hazmat.primitives.asymmetric import padding as asym_padding
+from cryptography.hazmat.primitives.asymmetric import rsa
+from rich.console import Console
+
+from ..config import Config
+
+app = typer.Typer(help="Login to Prime Intellect")
+console = Console()
+
+
+def generate_ephemeral_keypair() -> tuple[rsa.RSAPrivateKey, str]:
+    """Generate a temporary RSA key pair for secure communication"""
+    try:
+        private_key = rsa.generate_private_key(public_exponent=65537, key_size=2048)
+        public_key = private_key.public_key()
+
+        # Serialize public key to PEM format
+        public_pem = public_key.public_bytes(
+            encoding=serialization.Encoding.PEM,
+            format=serialization.PublicFormat.SubjectPublicKeyInfo,
+        ).decode("utf-8")
+
+        return private_key, public_pem
+    except Exception as e:
+        console.print(f"[red]Error generating keypair: {str(e)}[/red]")
+        raise typer.Exit(1)
+
+
+def decrypt_challenge_response(
+    private_key: rsa.RSAPrivateKey, encrypted_response: bytes
+) -> Optional[bytes]:
+    """Decrypt the challenge response using the private key"""
+    try:
+        decrypted: bytes = private_key.decrypt(
+            encrypted_response,
+            asym_padding.OAEP(
+                mgf=asym_padding.MGF1(algorithm=hashes.SHA256()),
+                algorithm=hashes.SHA256(),
+                label=None,
+            ),
+        )
+        return decrypted
+    except Exception as e:
+        console.print(f"[red]Error decrypting response: {str(e)}[/red]")
+        return None
+
+
+@app.callback(invoke_without_command=True)
+def login() -> None:
+    """Login to Prime Intellect"""
+    config = Config()
+    settings = config.view()
+
+    if not settings["base_url"]:
+        console.print(
+            "[red]Base URL not configured.",
+            "Please run 'prime config set-base-url' first.",
+        )
+        raise typer.Exit(1)
+
+    private_key = None
+    try:
+        # Generate secure keypair
+        private_key, public_pem = generate_ephemeral_keypair()
+
+        response = requests.post(
+            f"{settings['base_url']}/api/v1/auth_challenge/generate",
+            json={
+                "encryptionPublicKey": public_pem,
+            },
+        )
+
+        if response.status_code != 200:
+            console.print(
+                "[red]Failed to generate challenge:",
+                f"{response.json().get('detail', 'Unknown error')}[/red]",
+            )
+            raise typer.Exit(1)
+
+        challenge_response = response.json()
+
+        console.print("\n[bold blue]To login, please follow these steps:[/bold blue]")
+        console.print(
+            "1. Open ",
+            "[link]https://app.primeintellect.ai/dashboard/tokens/challenge[/link]",
+        )
+        console.print(
+            "2. Enter this code:",
+            f"[bold green]{challenge_response['challenge']}[/bold green]",
+        )
+        console.print("\nWaiting for authentication...")
+
+        challenge_auth_header = f"Bearer {challenge_response['status_auth_token']}"
+        while True:
+            try:
+                status_response = requests.get(
+                    f"{settings['base_url']}/api/v1/auth_challenge/status",
+                    params={"challenge": challenge_response["challenge"]},
+                    headers={"Authorization": challenge_auth_header},
+                )
+
+                if status_response.status_code == 404:
+                    console.print("[red]Challenge expired[/red]")
+                    break
+
+                status_data = status_response.json()
+                if status_data.get("result"):
+                    # Decrypt the result
+                    encrypted_result = base64.b64decode(status_data["result"])
+                    decrypted_result = decrypt_challenge_response(
+                        private_key, encrypted_result
+                    )
+                    if decrypted_result:
+                        # Update config with decrypted token
+                        config.set_api_key(decrypted_result.decode())
+                        console.print("[green]Successfully logged in![/green]")
+                    else:
+                        console.print(
+                            "[red]Failed to decrypt authentication token[/red]"
+                        )
+                    break
+
+                time.sleep(5)
+            except requests.exceptions.RequestException:
+                console.print("[red]Failed to connect to server. Retrying...[/red]")
+                time.sleep(5)
+                continue
+
+    except KeyboardInterrupt:
+        console.print("\n[yellow]Login cancelled by user[/yellow]")
+        raise typer.Exit(1)
+    except Exception as e:
+        console.print(f"[red]An error occurred: {str(e)}[/red]")
+        raise typer.Exit(1)
+    finally:
+        # Ensure private key is securely wiped
+        if private_key:
+            del private_key

src/prime_cli/commands/pods.py

@@ -2,6 +2,7 @@
 import json
 import os
 import subprocess
+import time
 from datetime import datetime
 from typing import Dict, List, Optional, Tuple, Union
 
@@ -752,16 +753,22 @@ def connect(pod_id: str) -> None:
         base_client = APIClient()
         pods_client = PodsClient(base_client)
 
-        # Get pod status to check SSH connection details
-        statuses = pods_client.get_status([pod_id])
-        if not statuses:
-            console.print(f"[red]No status found for pod {pod_id}[/red]")
-            raise typer.Exit(1)
+        # Keep trying until SSH connection is available
+        status_message = "[bold blue]Waiting for SSH connection to become available..."
+        with console.status(status_message, spinner="dots"):
+            while True:
+                # Get pod status to check SSH connection details
+                statuses = pods_client.get_status([pod_id])
+                if not statuses:
+                    console.print(f"[red]No status found for pod {pod_id}[/red]")
+                    raise typer.Exit(1)
 
-        status = statuses[0]
-        if not status.ssh_connection:
-            console.print(f"[red]SSH connection not available for pod {pod_id}[/red]")
-            raise typer.Exit(1)
+                status = statuses[0]
+                if status.ssh_connection:
+                    break
+
+                # Wait before retrying
+                time.sleep(5)  # Wait 5 seconds before retrying
 
         # Get SSH key path from config
         ssh_key_path = config.ssh_key_path

src/prime_cli/main.py

@@ -4,6 +4,7 @@
 
 from .commands.availability import app as availability_app
 from .commands.config import app as config_app
+from .commands.login import app as login_app
 from .commands.pods import app as pods_app
 
 __version__ = version("prime-cli")
@@ -13,6 +14,7 @@
 app.add_typer(availability_app, name="availability")
 app.add_typer(config_app, name="config")
 app.add_typer(pods_app, name="pods")
+app.add_typer(login_app, name="login")
 
 
 @app.callback(invoke_without_command=True)