gitarsenal-cli 1.1.2 → 1.1.4

package/python/modal_proxy_service.py

@@ -0,0 +1,336 @@
+#!/usr/bin/env python3
+"""
+Modal Proxy Service for GitArsenal CLI
+
+This service allows GitArsenal CLI users to access Modal services without exposing your Modal token.
+It acts as a secure proxy between clients and Modal's API.
+"""
+
+import os
+import json
+import secrets
+import string
+import time
+from flask import Flask, request, jsonify
+from flask_cors import CORS
+import modal
+import threading
+import logging
+from dotenv import load_dotenv
+import uuid
+import sys
+from pathlib import Path
+
+# Add the current directory to the path so we can import the test_modalSandboxScript module
+current_dir = Path(__file__).parent.absolute()
+sys.path.append(str(current_dir.parent))
+
+# Load environment variables from .env file
+load_dotenv()
+
+# Configure logging
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
+    handlers=[
+        logging.FileHandler("modal_proxy.log"),
+        logging.StreamHandler()
+    ]
+)
+logger = logging.getLogger("modal-proxy")
+
+app = Flask(__name__)
+CORS(app)  # Enable CORS for all routes
+
+# Get Modal token from environment variable
+MODAL_TOKEN = os.environ.get("MODAL_TOKEN")
+if not MODAL_TOKEN:
+    logger.error("MODAL_TOKEN environment variable is not set!")
+
+# Dictionary to store active containers
+active_containers = {}
+
+# Authentication tokens for clients
+# In a production environment, use a proper authentication system
+API_KEYS = {}
+if os.environ.get("API_KEYS"):
+    API_KEYS = {key.strip(): True for key in os.environ.get("API_KEYS").split(",")}
+
+def generate_api_key():
+    """Generate a new API key for a client"""
+    return ''.join(secrets.choice(string.ascii_letters + string.digits) for _ in range(32))
+
+def authenticate_request():
+    """Authenticate the request using API key"""
+    api_key = request.headers.get('X-API-Key')
+    if not api_key or api_key not in API_KEYS:
+        return False
+    return True
+
+def generate_random_password(length=16):
+    """Generate a random password for SSH access"""
+    alphabet = string.ascii_letters + string.digits + "!@#$%^&*"
+    password = ''.join(secrets.choice(alphabet) for i in range(length))
+    return password
+
+def setup_modal_auth():
+    """Set up Modal authentication using the server's token"""
+    if not MODAL_TOKEN:
+        logger.error("Cannot set up Modal authentication: No token provided")
+        return False
+    
+    try:
+        # Set the token in the environment
+        os.environ["MODAL_TOKEN_ID"] = MODAL_TOKEN
+        # Also set the token directly in modal.config
+        modal.config._auth_config.token_id = MODAL_TOKEN
+        logger.info("Modal token set in environment and config")
+        return True
+    except Exception as e:
+        logger.error(f"Error setting up Modal authentication: {e}")
+        return False
+
+@app.route('/api/health', methods=['GET'])
+def health_check():
+    """Health check endpoint"""
+    return jsonify({"status": "ok", "message": "Modal proxy service is running"})
+
+@app.route('/api/create-api-key', methods=['POST'])
+def create_api_key():
+    """Create a new API key (protected by admin key)"""
+    admin_key = request.headers.get('X-Admin-Key')
+    if not admin_key or admin_key != os.environ.get("ADMIN_KEY"):
+        return jsonify({"error": "Unauthorized"}), 401
+    
+    new_key = generate_api_key()
+    API_KEYS[new_key] = True
+    
+    return jsonify({"api_key": new_key})
+
+@app.route('/api/create-sandbox', methods=['POST'])
+def create_sandbox():
+    """Create a Modal sandbox"""
+    if not authenticate_request():
+        return jsonify({"error": "Unauthorized"}), 401
+    
+    if not setup_modal_auth():
+        return jsonify({"error": "Failed to set up Modal authentication"}), 500
+    
+    try:
+        data = request.json
+        gpu_type = data.get('gpu_type', 'A10G')
+        repo_url = data.get('repo_url')
+        repo_name = data.get('repo_name')
+        setup_commands = data.get('setup_commands', [])
+        volume_name = data.get('volume_name')
+        
+        # Import the sandbox creation function from your module
+        from test_modalSandboxScript import create_modal_sandbox
+        
+        logger.info(f"Creating sandbox with GPU: {gpu_type}, Repo: {repo_url}")
+        
+        # Create a unique ID for this sandbox
+        sandbox_id = str(uuid.uuid4())
+        
+        # Start sandbox creation in a separate thread
+        def create_sandbox_thread():
+            try:
+                # Ensure Modal token is set before creating sandbox
+                if not setup_modal_auth():
+                    logger.error("Failed to set up Modal authentication in thread")
+                    return
+                    
+                result = create_modal_sandbox(
+                    gpu_type, 
+                    repo_url=repo_url, 
+                    repo_name=repo_name, 
+                    setup_commands=setup_commands,
+                    volume_name=volume_name
+                )
+                
+                if result:
+                    active_containers[sandbox_id] = {
+                        "container_id": result.get("container_id"),
+                        "sandbox_id": result.get("sandbox_id"),
+                        "created_at": time.time(),
+                        "type": "sandbox"
+                    }
+                    logger.info(f"Sandbox created successfully: {result.get('container_id')}")
+                else:
+                    logger.error("Failed to create sandbox")
+            except Exception as e:
+                logger.error(f"Error in sandbox creation thread: {e}")
+        
+        thread = threading.Thread(target=create_sandbox_thread)
+        thread.daemon = True
+        thread.start()
+        
+        return jsonify({
+            "message": "Sandbox creation started",
+            "sandbox_id": sandbox_id
+        })
+        
+    except Exception as e:
+        logger.error(f"Error creating sandbox: {e}")
+        return jsonify({"error": str(e)}), 500
+
+@app.route('/api/create-ssh-container', methods=['POST'])
+def create_ssh_container():
+    """Create a Modal SSH container"""
+    if not authenticate_request():
+        return jsonify({"error": "Unauthorized"}), 401
+    
+    if not setup_modal_auth():
+        return jsonify({"error": "Failed to set up Modal authentication"}), 500
+    
+    try:
+        data = request.json
+        gpu_type = data.get('gpu_type', 'A10G')
+        repo_url = data.get('repo_url')
+        repo_name = data.get('repo_name')
+        setup_commands = data.get('setup_commands', [])
+        volume_name = data.get('volume_name')
+        timeout_minutes = data.get('timeout', 60)
+        
+        # Generate a random password for SSH
+        ssh_password = generate_random_password()
+        
+        # Import the SSH container creation function
+        from test_modalSandboxScript import create_modal_ssh_container
+        
+        logger.info(f"Creating SSH container with GPU: {gpu_type}, Repo: {repo_url}")
+        
+        # Create a unique ID for this container
+        container_id = str(uuid.uuid4())
+        
+        # Start container creation in a separate thread
+        def create_container_thread():
+            try:
+                # Ensure Modal token is set before creating container
+                if not setup_modal_auth():
+                    logger.error("Failed to set up Modal authentication in thread")
+                    return
+                
+                # Explicitly set the Modal token in the environment again for this thread
+                os.environ["MODAL_TOKEN_ID"] = MODAL_TOKEN
+                
+                # Log token status for debugging
+                token_status = "Token is set" if MODAL_TOKEN else "Token is missing"
+                logger.info(f"Modal token status: {token_status}")
+                
+                result = create_modal_ssh_container(
+                    gpu_type, 
+                    repo_url=repo_url, 
+                    repo_name=repo_name, 
+                    setup_commands=setup_commands,
+                    volume_name=volume_name,
+                    timeout_minutes=timeout_minutes,
+                    ssh_password=ssh_password
+                )
+                
+                if result:
+                    active_containers[container_id] = {
+                        "container_id": result.get("app_name"),
+                        "ssh_password": ssh_password,
+                        "created_at": time.time(),
+                        "type": "ssh"
+                    }
+                    logger.info(f"SSH container created successfully: {result.get('app_name')}")
+                else:
+                    logger.error("Failed to create SSH container")
+            except Exception as e:
+                logger.error(f"Error in SSH container creation thread: {e}")
+        
+        thread = threading.Thread(target=create_container_thread)
+        thread.daemon = True
+        thread.start()
+        
+        return jsonify({
+            "message": "SSH container creation started",
+            "container_id": container_id,
+            "ssh_password": ssh_password
+        })
+        
+    except Exception as e:
+        logger.error(f"Error creating SSH container: {e}")
+        return jsonify({"error": str(e)}), 500
+
+@app.route('/api/container-status/<container_id>', methods=['GET'])
+def container_status(container_id):
+    """Get status of a container"""
+    if not authenticate_request():
+        return jsonify({"error": "Unauthorized"}), 401
+    
+    if container_id in active_containers:
+        # Remove sensitive information like passwords
+        container_info = active_containers[container_id].copy()
+        if "ssh_password" in container_info:
+            del container_info["ssh_password"]
+        
+        return jsonify({
+            "status": "active",
+            "info": container_info
+        })
+    else:
+        return jsonify({
+            "status": "not_found",
+            "message": "Container not found or has been terminated"
+        }), 404
+
+@app.route('/api/terminate-container', methods=['POST'])
+def terminate_container():
+    """Terminate a Modal container"""
+    if not authenticate_request():
+        return jsonify({"error": "Unauthorized"}), 401
+    
+    if not setup_modal_auth():
+        return jsonify({"error": "Failed to set up Modal authentication"}), 500
+    
+    try:
+        data = request.json
+        container_id = data.get('container_id')
+        
+        if not container_id:
+            return jsonify({"error": "Container ID is required"}), 400
+        
+        if container_id not in active_containers:
+            return jsonify({"error": "Container not found"}), 404
+        
+        modal_container_id = active_containers[container_id].get("container_id")
+        
+        # Terminate the container using Modal CLI
+        import subprocess
+        result = subprocess.run(
+            ["modal", "container", "terminate", modal_container_id],
+            capture_output=True,
+            text=True
+        )
+        
+        if result.returncode == 0:
+            # Remove from active containers
+            del active_containers[container_id]
+            logger.info(f"Container terminated successfully: {modal_container_id}")
+            return jsonify({"message": "Container terminated successfully"})
+        else:
+            logger.error(f"Failed to terminate container: {result.stderr}")
+            return jsonify({"error": f"Failed to terminate container: {result.stderr}"}), 500
+        
+    except Exception as e:
+        logger.error(f"Error terminating container: {e}")
+        return jsonify({"error": str(e)}), 500
+
+if __name__ == '__main__':
+    # Check if Modal token is set
+    if not MODAL_TOKEN:
+        logger.error("MODAL_TOKEN environment variable must be set!")
+        exit(1)
+    
+    # Generate an admin key if not set
+    if not os.environ.get("ADMIN_KEY"):
+        admin_key = generate_api_key()
+        os.environ["ADMIN_KEY"] = admin_key
+        logger.info(f"Generated admin key: {admin_key}")
+        print(f"Admin key: {admin_key}")
+    
+    port = int(os.environ.get("PORT", 5001))  # Default to 5001 to avoid macOS AirPlay conflict
+    app.run(host='0.0.0.0', port=port) 

package/python/requirements.txt

@@ -1,4 +1,6 @@
 modal>=0.56.4
 requests>=2.31.0
 pathlib>=1.0.1
-python-dotenv>=1.0.0 
+python-dotenv>=1.0.0
+flask>=2.0.0
+flask-cors>=3.0.0 

package/python/test_modalSandboxScript.py

@@ -2067,6 +2067,38 @@ def create_modal_ssh_container(gpu_type, repo_url=None, repo_name=None, setup_co
                                volume_name=None, timeout_minutes=60, ssh_password=None):
     """Create a Modal SSH container with GPU support and tunneling"""
     
+    # Check if Modal is authenticated
+    try:
+        # Try to access Modal token to check authentication
+        try:
+            # This will raise an exception if not authenticated
+            modal.config.get_current_workspace_name()
+            print("✅ Modal authentication verified")
+        except modal.exception.AuthError:
+            print("\n" + "="*80)
+            print("🔑 MODAL AUTHENTICATION REQUIRED")
+            print("="*80)
+            print("GitArsenal requires Modal authentication to create cloud environments.")
+            
+            # Check if token is in environment
+            modal_token = os.environ.get("MODAL_TOKEN_ID")
+            if not modal_token:
+                print("⚠️ No Modal token found in environment.")
+                return None
+            else:
+                print("🔄 Using Modal token from environment")
+                # Try to authenticate with the token
+                try:
+                    # Set token directly in modal config
+                    modal.config._auth_config.token_id = modal_token
+                    print("✅ Modal token set in config")
+                except Exception as e:
+                    print(f"⚠️ Error setting Modal token: {e}")
+                    return None
+    except Exception as e:
+        print(f"⚠️ Error checking Modal authentication: {e}")
+        print("Continuing anyway, but Modal operations may fail")
+    
     # Generate a unique app name with timestamp to avoid conflicts
     timestamp = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
     app_name = f"ssh-container-{timestamp}"
@@ -2116,6 +2148,18 @@ def create_modal_ssh_container(gpu_type, repo_url=None, repo_name=None, setup_co
             print(f"⚠️ Could not create default volume: {e}")
             print("⚠️ Continuing without persistent volume")
             volume = None
+    
+    # Print debug info for authentication
+    print("🔍 Modal authentication debug info:")
+    modal_token = os.environ.get("MODAL_TOKEN_ID")
+    print(f"  - MODAL_TOKEN_ID in env: {'Yes' if modal_token else 'No'}")
+    print(f"  - Token length: {len(modal_token) if modal_token else 'N/A'}")
+    
+    try:
+        workspace = modal.config.get_current_workspace_name()
+        print(f"  - Current workspace: {workspace}")
+    except Exception as e:
+        print(f"  - Error getting workspace: {e}")
 
     # Create SSH-enabled image
     ssh_image = (
@@ -2153,7 +2197,12 @@ def create_modal_ssh_container(gpu_type, repo_url=None, repo_name=None, setup_co
     )
 
     # Create the Modal app
-    app = modal.App(app_name)
+    try:
+        app = modal.App(app_name)
+        print("✅ Created Modal app successfully")
+    except Exception as e:
+        print(f"❌ Error creating Modal app: {e}")
+        return None
     
     # Configure volumes if available
     volumes_config = {}
@@ -2319,6 +2368,10 @@ def create_modal_ssh_container(gpu_type, repo_url=None, repo_name=None, setup_co
             "volume_name": volume_name,
             "volume_mount_path": volume_mount_path if volume else None
         }
+    except modal.exception.AuthError as auth_err:
+        print(f"❌ Modal authentication error: {auth_err}")
+        print("🔑 Please check that your Modal token is valid and properly set")
+        return None
     except KeyboardInterrupt:
         print("\n👋 Container startup interrupted")
         return None