tfm_ainventory/backend/main.py

import os
from ipaddress import ip_address, ip_network, AddressValueError
from .config_loader import get_config
from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from slowapi import Limiter
from slowapi.util import get_remote_address
from . import models
from .database import engine
from .routers import items, operations, users, auth, sync, categories
from .routers.admin import backups, ai_config, db_config, exports
from .logger import log
from .scheduler import scheduler, sync_scheduler_config
from .services.image_storage import ensure_image_directories, IMAGES_ROOT

# Create the database tables
from .database import DATA_DIR, db_path
log.info(f"Using DATA_DIR: {DATA_DIR}")
log.info(f"Database path: {db_path}")
models.Base.metadata.create_all(bind=engine)
log.info("Database tables verified.")

app = FastAPI(title="TFM aInventory API", version="1.1.0")
log.info("TFM aInventory API process started.")

# [SECURITY FIX M-01] CORS Configuration with Subnet Support
# We dynamically build allowed origins from configuration to simplify deployment.
config = get_config()
app_config = config.get("application", {})
_raw_origins = app_config.get("cors_origins", "")
ALLOWED_ORIGINS = [o.strip() for o in _raw_origins.split(",") if o.strip()]

# Allowed subnets for subnet-based CORS validation (e.g., VPN, Tailscale)
ALLOWED_SUBNETS = []

# Automatically add origins based on network config if present
server_ip = app_config.get("server_ip")
front_port = app_config.get("frontend_port", "8917")
front_ssl_port = app_config.get("frontend_ssl_port", "8919")
back_ssl_port = app_config.get("backend_ssl_port", "8918")

# Always allow localhost
defaults = [
    f"http://localhost:{front_port}",
    f"https://localhost:{front_ssl_port}",
    f"https://localhost:{back_ssl_port}",
]
for d in defaults:
    if d not in ALLOWED_ORIGINS:
        ALLOWED_ORIGINS.append(d)

# Add IP-based origins if SERVER_IP is set
if server_ip and server_ip != "localhost":
    ip_origins = [
        f"http://{server_ip}:{front_port}",
        f"https://{server_ip}:{front_ssl_port}",
        f"https://{server_ip}:{back_ssl_port}",
    ]
    for ip_o in ip_origins:
        if ip_o not in ALLOWED_ORIGINS:
            ALLOWED_ORIGINS.append(ip_o)

# [NEW] Add Extra Allowed Origins (Tailscale, VPN, etc.) with Subnet Support
extra_origins_raw = app_config.get("cors_origins", "")
if extra_origins_raw:
    for extra_item in [o.strip() for o in extra_origins_raw.split(",") if o.strip()]:
        # Check if it's a subnet (contains /) or individual IP
        if "/" in extra_item:
            try:
                # Parse as subnet
                subnet = ip_network(extra_item, strict=False)
                ALLOWED_SUBNETS.append(subnet)
                log.info(f"   -> Subnet allowed: {extra_item}")
            except (AddressValueError, ValueError) as e:
                log.warning(f"   ⚠️  Invalid subnet {extra_item}: {e}")
        else:
            # Treat as individual IP - generate standard port combinations
            ext_combos = [
                f"http://{extra_item}:{front_port}",
                f"https://{extra_item}:{front_ssl_port}",
                f"https://{extra_item}:{back_ssl_port}",
            ]
            for combo in ext_combos:
                if combo not in ALLOWED_ORIGINS:
                    ALLOWED_ORIGINS.append(combo)

log.info("🔒 [SECURITY] CORS configuration initialized.")
log.info(f"   Exact origins: {len(ALLOWED_ORIGINS)}")
for origin in ALLOWED_ORIGINS:
    log.info(f"      -> {origin}")
if ALLOWED_SUBNETS:
    log.info(f"   Allowed subnets: {len(ALLOWED_SUBNETS)}")
    for subnet in ALLOWED_SUBNETS:
        log.info(f"      -> {subnet}")

# Helper function to check if origin is allowed (exact match or subnet)
def is_origin_allowed(origin: str) -> bool:
    """Check if origin is in allowed origins or matches any allowed subnet"""
    # Check exact match first (faster)
    if origin in ALLOWED_ORIGINS:
        return True

    # Check subnet match if subnets are configured
    if not ALLOWED_SUBNETS:
        return False

    try:
        # Extract IP from origin URL (e.g., "https://192.168.1.100:8919" -> "192.168.1.100")
        from urllib.parse import urlparse
        parsed = urlparse(origin)
        origin_host = parsed.hostname
        if not origin_host:
            return False

        origin_ip = ip_address(origin_host)
        for subnet in ALLOWED_SUBNETS:
            if origin_ip in subnet:
                return True
    except (AddressValueError, ValueError):
        pass

    return False

# Add CORS middleware FIRST (before rate limiter)
# Uses is_origin_allowed() to validate exact origins + subnet matching
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.requests import Request
from starlette.responses import Response

class SubnetAwareCORSMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request: Request, call_next) -> Response:
        origin = request.headers.get("origin")

        # Handle CORS preflight (OPTIONS) requests FIRST
        if request.method == "OPTIONS":
            if origin and is_origin_allowed(origin):
                return Response(
                    status_code=200,
                    headers={
                        "Access-Control-Allow-Origin": origin,
                        "Access-Control-Allow-Credentials": "true",
                        "Access-Control-Allow-Methods": "GET, POST, PUT, PATCH, DELETE, OPTIONS",
                        "Access-Control-Allow-Headers": "*",
                        "Content-Length": "0",
                    }
                )
            return Response(status_code=403)

        # Process the actual request
        response = await call_next(request)

        # Add CORS headers to response if origin is allowed
        if origin and is_origin_allowed(origin):
            response.headers["Access-Control-Allow-Origin"] = origin
            response.headers["Access-Control-Allow-Credentials"] = "true"
            response.headers["Access-Control-Allow-Methods"] = "GET, POST, PUT, PATCH, DELETE, OPTIONS"
            response.headers["Access-Control-Allow-Headers"] = "*"

        return response

app.add_middleware(SubnetAwareCORSMiddleware)
log.info("🔒 [CORS] Subnet-aware middleware enabled (exact origins + subnet matching)")

# [H-02] Rate limiting on API
limiter = Limiter(key_func=get_remote_address)
app.state.limiter = limiter

app.include_router(items.router)
app.include_router(operations.router)
app.include_router(users.router)
app.include_router(auth.router)
app.include_router(sync.router)
app.include_router(categories.router)
app.include_router(backups.router)
app.include_router(ai_config.router)
app.include_router(db_config.router)
app.include_router(exports.router)

# [STATIC FILES] Mount /images/ directory for serving uploaded photos
# Ensure directory exists before mounting (StaticFiles requires pre-existing directory)
IMAGES_ROOT.mkdir(parents=True, exist_ok=True)
# Mount at root level after all API routes to avoid conflicts with dynamic routes.
# MIME types are auto-detected by StaticFiles based on file extensions.
# Supported formats: JPEG (.jpg/.jpeg), PNG (.png), WebP (.webp), GIF (.gif)
app.mount("/images", StaticFiles(directory=str(IMAGES_ROOT)), name="images")

@app.on_event("startup")
def startup_event():
    log.info("[STARTUP] Initializing image storage directories...")
    ensure_image_directories()
    log.info("[STARTUP] Starting background scheduler...")
    scheduler.start()
    sync_scheduler_config()

    # [NEW] Initialize default system settings
    from .database import SessionLocal
    db = SessionLocal()
    try:
        # Default AI Prompt from User Request
        default_prompt = (
            "Extract hardware specifications with PRECISE formatting.\n"
            "For any field, do not add comments in parenthesis or measurement units in Item name.\n\n"
            "SIZE CONVERSION RULES (CRITICAL - Apply to all size fields):\n"
            "  Convert sizes to HUMAN-READABLE format:\n"
            "  • Storage ≥1600GB: convert to TB (e.g., '1600GB' → '1.6TB', '2048GB' → '2TB', '8192GB' → '8TB')\n"
            "  • Storage 512GB-1599GB: keep as GB (e.g., '512GB', '1TB' if labeled as 1000GB)\n"
            "  • Memory ≥1024MB: convert to GB (e.g., '1024MB' → '1GB')\n"
            "  • Cables/length: express in meters (e.g., '5m', '30m')\n"
            "  Apply this rule to: Item field size prefix, Size field, and OCR key sizes.\n\n"
            "ITEM FIELD FORMAT (Critical):\n"
            "[<size_or_length>] <type> <vendor> <connector> <part_number>\n"
            "  <size_or_length>: HUMAN-READABLE format - GB/TB for storage, meters for cables (e.g., '2m', '256GB', '1.6TB'). Omit diameter/mm measurements.\n"
            "  <type>: DDR3/DDR4/DDR5/SSD/HDD/NVMe/SAS/SATA/Patchcord/Fiber/Cable/Transceiver etc.\n"
            "  <vendor>: HP/HPE/Dell/Samsung/Cisco/Lenovo/Hynix etc.\n"
            "  <connector>: RJ45/LC-LC/MPO/U.3/SATA/SAS/LC/ST etc. Omit if N/A.\n"
            "  <part_number>: PN only if visible on label. Omit serial numbers.\n"
            "Examples: '5m Patchcord LC-LC' / '256GB SSD Samsung SAS' / '128GB DDR4 Hynix SK-234' / '1.6TB NVMe HP U.3' / '8TB HDD Seagate SATA'\n\n"
            "TYPE: Item asset class (DDR3/SSD/NVMe/Patchcord/SFP etc.)\n"
            "DESCRIPTION: Technical details max 5 words (e.g., 'High speed fiber optic cable'). Omit size/length here.\n"
            "CATEGORY: Broad ecosystem (Memory, Storage, Network, Cabling, etc.)\n"
            "CONNECTOR: Physical interface type (e.g., 'LC', 'RJ45', 'U.3')\n"
            "SIZE: HUMAN-READABLE capacity or length ONLY (e.g., '256GB', '1.6TB', '2m'). No serial numbers or comments.\n"
            "COLOR: Physical color if distinguishing.\n"
            "PartNr: Part number only (no serial numbers)\n"
            "OCR: Robust matching key. Include: core type + vendor + key identifiers + HUMAN-READABLE size + common variations.\n"
            "     Examples: 'PATCHCORD 5M LC LC CAT6 FIBER' / '256GB SSD SAMSUNG' / '1.6TB NVME HP' / '128GB DDR4 HYNIX'\n"
            "     Include abbreviations: 'DDR4'/'DDR' or 'SSD'/'SSDS' or 'HP'/'HPE' or 'NVMe'/'NVME'\n"
            "     Format: ALL UPPERCASE, space-separated tokens, NO special chars. Skip serial numbers. Use human-readable sizes."
        )
        
        # Wrap in JSON instructions for reliable parsing
        final_prompt = f"IMAGE ANALYSIS INSTRUCTIONS:\n{default_prompt}\n\nIMPORTANT: Return ONLY a valid JSON object with the keys: Item, Type, Description, Category, Connector, Size, Color, PartNr, OCR."

        existing = db.query(models.SystemSetting).filter(models.SystemSetting.key == "ai_extraction_prompt").first()
        if not existing:
            db.add(models.SystemSetting(key="ai_extraction_prompt", value=final_prompt))
            db.commit()
            log.info("Initialized default AI prompt in database.")
            # Refresh existing after commit
            existing = db.query(models.SystemSetting).filter(models.SystemSetting.key == "ai_extraction_prompt").first()

        # [NEW] Sync/Initialize AI prompt file in /config
        from .database import BASE_DIR
        PROJECT_ROOT = os.path.dirname(BASE_DIR)
        PROMPT_FILE_PATH = os.path.join(PROJECT_ROOT, "config", "ai_prompt.md")
        
        if not os.path.exists(PROMPT_FILE_PATH):
            try:
                os.makedirs(os.path.dirname(PROMPT_FILE_PATH), exist_ok=True)
                # Use DB value (which we just ensured exists)
                current_val = existing.value if existing else final_prompt
                with open(PROMPT_FILE_PATH, 'w', encoding='utf-8') as f:
                    f.write(current_val)
                log.info(f"✅ Initialized AI prompt configuration file: {PROMPT_FILE_PATH}")
            except Exception as fe:
                log.error(f"❌ Failed to initialize AI prompt file: {fe}")

        defaults = {
            "backup_retention_count": "10",
            "backup_schedule_hour": "3",
            "backup_schedule_freq_days": "1",
            "ai_provider": "gemini"
        }
        for key, val in defaults.items():
            if not db.query(models.SystemSetting).filter(models.SystemSetting.key == key).first():
                db.add(models.SystemSetting(key=key, value=val))
                log.info(f"Initialized default setting: {key}")
        db.commit()
    except Exception as e:
        log.error(f"Failed to initialize settings: {e}")
    finally:
        db.close()

@app.get("/")
def read_root():
    return {"message": "Inventory API is running"}