main.py

import os
from contextlib import asynccontextmanager
from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from pathlib import Path
import asyncio

# Import middleware and utilities
# from middleware.tracking import TrackingMiddleware
from utils.sliver_client import get_client, close_client
from database import init_db
from utils.geoip import geoip_manager

@asynccontextmanager
async def lifespan(app: FastAPI):
    """Manage the application lifecycle"""
    # Startup code
    await init_db()

    yield
    # Shutdown code
    await close_client()
    if geoip_manager:
        geoip_manager.close()

# Create FastAPI app
app = FastAPI(
    title="Sliver API",
    description="REST API for interacting with Sliver C2",
    version="1.0.0",
    lifespan=lifespan
)

# Add CORS middleware
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
    expose_headers=["*"],
    max_age=3600,
)

# Add tracking middleware
# tracking_middleware = TrackingMiddleware(geoip_manager)
# app.add_middleware(type(tracking_middleware), instance=tracking_middleware)

# Import route modules
from routes.auth import router as auth_router
from routes.users import router as users_router
from routes.connections import router as connections_router
from routes.sessions import router as sessions_router
from routes.listeners import router as listeners_router
from routes.processes import router as processes_router
from routes.network import router as network_router
from routes.environment import router as environment_router
from routes.registry import router as registry_router
from routes.implants import router as implants_router
from routes.sliver import router as sliver_router
from routes.uploader_routes import router as uploader_router
from routes.initial_access_vectors import router as intiial_access_vector_router
from routes.beacons import router as beacons_router

# Include routers
app.include_router(auth_router, prefix="/auth", tags=["authentication"])
app.include_router(users_router, prefix="/users", tags=["users"])
app.include_router(connections_router, prefix="/connections", tags=["connections"])
app.include_router(sessions_router, prefix="/sessions")
app.include_router(listeners_router, prefix="/listeners", tags=["listeners"])
app.include_router(processes_router, prefix="/processes", tags=["processes"])
app.include_router(network_router, prefix="/network", tags=["network"])
app.include_router(environment_router, prefix="/environment", tags=["environment"])
app.include_router(registry_router, prefix="/registry", tags=["registry"])
app.include_router(implants_router, prefix="/implant", tags=["implants"])
app.include_router(intiial_access_vector_router, tags=["initial_access_vectors"])
app.include_router(sliver_router, tags=["sliver"])
app.include_router(uploader_router, tags=["uploadroutes"]) # POST /upload, # GET /map
app.include_router(beacons_router, prefix="/beacons", tags=["beacons"])

@app.get("/health")
async def health_check():
    """Service health check"""
    return {"status": "healthy"}

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)