from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse
from pydantic import BaseModel
import subprocess
import shlex
import os
from fastapi import FastAPI, HTTPException, Request
app = FastAPI()

# Add CORS middleware
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # Allows all origins
    allow_credentials=True,
    allow_methods=["*"],  # Allows all methods
    allow_headers=["*"],  # Allows all headers
)

# Mount the static files directory
app.mount("/static", StaticFiles(directory="static"), name="static")

ALLOWED_COMMANDS = {
    'ls', 'cd', 'pwd', 'echo', 'cat', 'grep', 'find', 'touch', 'mkdir', 'rm', 'cp', 'mv'
}

class Command(BaseModel):
    command: str

@app.post("/execute")
async def execute_command(command: Command, request: Request):
    try:
        # Get or create a session-specific directory
        session = request.session
        if 'working_dir' not in session:
            session['working_dir'] = '/tmp/user_' + os.urandom(8).hex()
            os.makedirs(session['working_dir'], exist_ok=True)
        
        working_dir = session['working_dir']
        
        # Parse the command to get the base command
        base_command = shlex.split(command.command)[0]
        
        # Check if the base command is allowed
        if base_command not in ALLOWED_COMMANDS:
            raise HTTPException(status_code=403, detail=f"Command '{base_command}' is not allowed")

        # Handle 'cd' command specially
        if base_command == 'cd':
            new_dir = os.path.join(working_dir, shlex.split(command.command)[1])
            if os.path.isdir(new_dir):
                session['working_dir'] = new_dir
                return {"output": "", "error": "", "returncode": 0, "currentDirectory": new_dir}
            else:
                return {"output": "", "error": "No such directory", "returncode": 1, "currentDirectory": working_dir}

        # Execute the command in the controlled environment
        result = subprocess.run(command.command, shell=True, capture_output=True, text=True, timeout=5, cwd=working_dir)
        return {
            "output": result.stdout,
            "error": result.stderr,
            "returncode": result.returncode,
            "currentDirectory": working_dir
        }
    except subprocess.TimeoutExpired:
        raise HTTPException(status_code=408, detail="Command execution timed out")
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/")
async def root():
    return FileResponse('static/index.html')