firegex/backend/routers/nfregex.py

from base64 import b64decode
import secrets
import sqlite3
from fastapi import APIRouter, HTTPException
from fastapi.responses import PlainTextResponse
from pydantic import BaseModel
from modules.nfregex.nftables import FiregexTables
from modules.nfregex.firewall import STATUS, FirewallManager
from utils.sqlite import SQLite
from utils import ip_parse, refactor_name, socketio_emit, PortType
from modules.tls.service import activate_stream
from utils.models import ResetRequest, StatusMessageModel
from modules.nfregex.firegex import test_regex_validity

class ServiceModel(BaseModel):
    status: str
    service_id: str
    port: PortType|None = None
    name: str
    proto: str
    ip_int: str|None = None
    n_regex: int
    n_packets: int
    fail_open: bool
    target_type: str
    tls_stream_id: str|None = None

class RenameForm(BaseModel):
    name:str

class SettingsForm(BaseModel):
    port: PortType|None = None
    proto: str|None = None
    ip_int: str|None = None
    fail_open: bool|None = None
    target_type: str|None = None
    tls_stream_id: str|None = None

class RegexModel(BaseModel):
    regex:str
    mode:str
    id:int
    service_id:str
    n_packets:int
    is_case_sensitive:bool
    active:bool

class RegexAddForm(BaseModel):
    service_id: str
    regex: str
    mode: str
    active: bool|None = None
    is_case_sensitive: bool

class ServiceAddForm(BaseModel):
    name: str
    port: PortType|None = None
    proto: str
    ip_int: str|None = None
    fail_open: bool = False
    target_type: str = "flow"
    tls_stream_id: str|None = None

class ServiceAddResponse(BaseModel):
    status:str
    service_id: str|None = None

app = APIRouter()

db = SQLite('db/nft-regex.db', {
    'services': {
        'service_id': 'VARCHAR(100) PRIMARY KEY',
        'status': 'VARCHAR(100) NOT NULL',
        'target_type': 'VARCHAR(10) NOT NULL CHECK(target_type IN ("flow", "tls")) DEFAULT "flow"',
        'tls_stream_id': 'VARCHAR(100)',
        'port': 'INT CHECK(port > 0 and port < 65536)',
        'name': 'VARCHAR(100) NOT NULL UNIQUE',
        'proto': 'VARCHAR(3) NOT NULL CHECK (proto IN ("tcp", "udp"))',
        'ip_int': 'VARCHAR(100)',
        'fail_open': 'BOOLEAN NOT NULL CHECK (fail_open IN (0, 1)) DEFAULT 1',
    },
    'regexes': {
        'regex': 'TEXT NOT NULL',
        'mode': 'VARCHAR(1) NOT NULL CHECK (mode IN ("C", "S", "B"))', # C = to the client, S = to the server, B = both
        'service_id': 'VARCHAR(100) NOT NULL',
        'blocked_packets': 'INTEGER UNSIGNED NOT NULL DEFAULT 0',
        'regex_id': 'INTEGER PRIMARY KEY',
        'is_case_sensitive' : 'BOOLEAN NOT NULL CHECK (is_case_sensitive IN (0, 1))',
        'active' : 'BOOLEAN NOT NULL CHECK (active IN (0, 1)) DEFAULT 1',
        'FOREIGN KEY (service_id)':'REFERENCES services (service_id)',
    },
    'QUERY':[
        "CREATE UNIQUE INDEX IF NOT EXISTS unique_services ON services (port, ip_int, proto);",
        "CREATE UNIQUE INDEX IF NOT EXISTS unique_regex_service ON regexes (regex,service_id,mode,is_case_sensitive);"   
    ]
})

async def refresh_frontend(additional:list[str]=[]):
    await socketio_emit(["nfregex"]+additional)

async def reset(params: ResetRequest):
    if not params.delete: 
        db.backup()
    await firewall.close()
    FiregexTables().reset()
    if params.delete:
        db.delete()
        db.init()
    else:
        db.restore()
    try:
        await firewall.init()
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))
    

async def startup():
    db.init()
    try:
        await firewall.init()
    except Exception as e:
        print("WARNING cannot start firewall:", e)

async def shutdown():
    await firewall.close()
    db.disconnect()

def gen_service_id():
    while True:
        res = secrets.token_hex(8)
        if len(db.query('SELECT 1 FROM services WHERE service_id = ?;', res)) == 0:
            break
    return res

firewall = FirewallManager(db)

@app.get('/services', response_model=list[ServiceModel])
async def get_service_list():
    """Get the list of existent firegex services"""
    res = db.query("""
        SELECT
            s.service_id service_id,
            s.status status,
            s.port port,
            s.name name,
            s.proto proto,
            s.ip_int ip_int,
            s.fail_open fail_open,
            s.target_type target_type,
            s.tls_stream_id tls_stream_id,
            COUNT(r.regex_id) n_regex,
            COALESCE(SUM(r.blocked_packets),0) n_packets
        FROM services s LEFT JOIN regexes r ON s.service_id = r.service_id
        GROUP BY s.service_id;
    """)
    return res

@app.get('/services/{service_id}', response_model=ServiceModel)
async def get_service_by_id(service_id: str):
    """Get info about a specific service using his id"""
    res = db.query("""
        SELECT 
            s.service_id service_id,
            s.status status,
            s.port port,
            s.name name,
            s.proto proto,
            s.ip_int ip_int,
            s.fail_open fail_open,
            s.target_type target_type,
            s.tls_stream_id tls_stream_id,
            COUNT(r.regex_id) n_regex,
            COALESCE(SUM(r.blocked_packets),0) n_packets
        FROM services s LEFT JOIN regexes r ON s.service_id = r.service_id
        WHERE s.service_id = ? GROUP BY s.service_id;
    """, service_id)
    if len(res) == 0:
        raise HTTPException(status_code=400, detail="This service does not exists!")
    return res[0]

@app.post('/services/{service_id}/stop', response_model=StatusMessageModel)
async def service_stop(service_id: str):
    """Request the stop of a specific service"""
    await firewall.get(service_id).next(STATUS.STOP)
    await refresh_frontend()
    return {'status': 'ok'}

@app.post('/services/{service_id}/start', response_model=StatusMessageModel)
async def service_start(service_id: str):
    """Request the start of a specific service"""
    srv = db.query("SELECT target_type, tls_stream_id FROM services WHERE service_id = ?;", service_id)
    if srv and srv[0]["target_type"] == "tls":
        if not await activate_stream(srv[0]["tls_stream_id"]):
            raise HTTPException(status_code=400, detail="Linked TLS stream not found")
    await firewall.get(service_id).next(STATUS.ACTIVE)
    await refresh_frontend()
    return {'status': 'ok'}

@app.delete('/services/{service_id}', response_model=StatusMessageModel)
async def service_delete(service_id: str):
    """Request the deletion of a specific service"""
    db.query('DELETE FROM services WHERE service_id = ?;', service_id)
    db.query('DELETE FROM regexes WHERE service_id = ?;', service_id)
    await firewall.remove(service_id)
    await refresh_frontend()
    return {'status': 'ok'}

@app.put('/services/{service_id}/rename', response_model=StatusMessageModel)
async def service_rename(service_id: str, form: RenameForm):
    """Request to change the name of a specific service"""
    form.name = refactor_name(form.name)
    if not form.name:
        raise HTTPException(status_code=400, detail="The name cannot be empty!") 
    try:
        db.query('UPDATE services SET name=? WHERE service_id = ?;', form.name, service_id)
    except sqlite3.IntegrityError:
        raise HTTPException(status_code=400, detail="This name is already used")
    await refresh_frontend()
    return {'status': 'ok'}

@app.put('/services/{service_id}/settings', response_model=StatusMessageModel)
async def service_settings(service_id: str, form: SettingsForm):
    """Request to change the settings of a specific service (will cause a restart)"""
        
    srv_check = db.query("SELECT proto, ip_int, port, target_type, tls_stream_id FROM services WHERE service_id = ?;", service_id)
    if len(srv_check) == 0:
        raise HTTPException(status_code=404, detail="Service not found")
    old_srv = srv_check[0]
    
    if form.proto is not None and form.proto not in ["tcp", "udp"]:
        raise HTTPException(status_code=400, detail="Invalid protocol")
        
    if form.port is not None and (form.port < 1 or form.port > 65535):
        raise HTTPException(status_code=400, detail="Invalid port")
    
    if form.ip_int is not None:
        try:
            form.ip_int = ip_parse(form.ip_int)
        except ValueError:
            raise HTTPException(status_code=400, detail="Invalid address")
            
    new_ip = form.ip_int if form.ip_int is not None else old_srv["ip_int"]
    new_port = form.port if form.port is not None else old_srv["port"]
    new_proto = form.proto if form.proto is not None else old_srv["proto"]
    
    new_target_type = form.target_type if form.target_type is not None else old_srv["target_type"]
    new_tls_stream_id = form.tls_stream_id if form.tls_stream_id is not None else old_srv["tls_stream_id"]
    
    if new_target_type == "tls":
        if new_proto != "tcp":
            raise HTTPException(status_code=400, detail="TLS is only supported for TCP protocol")
        if not new_tls_stream_id:
            raise HTTPException(status_code=400, detail="TLS stream ID is required when target type is tls")
            
    if new_target_type == "flow" and (new_ip is None or new_port is None):
        raise HTTPException(status_code=400, detail="IP and Port are required when target type is flow")
    
    keys = []
    values = []
    
    for key, value in form.model_dump(exclude_none=True).items():
        keys.append(key)
        values.append(value)
        
    if len(keys) == 0:
        raise HTTPException(status_code=400, detail="No settings to change provided")
            
    try:
        db.query(f'UPDATE services SET {", ".join([f"{key}=?" for key in keys])} WHERE service_id = ?;', *values, service_id)
    except sqlite3.IntegrityError:
        raise HTTPException(status_code=400, detail="A service with these settings already exists")
    
    old_status = firewall.get(service_id).status
    await firewall.remove(service_id)
    if old_status == STATUS.ACTIVE and new_target_type == "tls":
        await activate_stream(new_tls_stream_id)
    await firewall.reload()
    await firewall.get(service_id).next(old_status)

    await refresh_frontend()
    return {'status': 'ok'}

@app.get('/services/{service_id}/regexes', response_model=list[RegexModel])
async def get_service_regexe_list(service_id: str):
    """Get the list of the regexes of a service"""
    if not db.query("SELECT 1 FROM services s WHERE s.service_id = ?;", service_id):
        raise HTTPException(status_code=400, detail="This service does not exists!")
    return db.query("""
        SELECT 
            regex, mode, regex_id `id`, service_id,
            blocked_packets n_packets, is_case_sensitive, active
        FROM regexes WHERE service_id = ?;
    """, service_id)

@app.get('/regexes/{regex_id}', response_model=RegexModel)
async def get_regex_by_id(regex_id: int):
    """Get regex info using his id"""
    res = db.query("""
        SELECT 
            regex, mode, regex_id `id`, service_id,
            blocked_packets n_packets, is_case_sensitive, active
        FROM regexes WHERE `id` = ?;
    """, regex_id)
    if len(res) == 0:
        raise HTTPException(status_code=400, detail="This regex does not exists!")
    return res[0]

@app.delete('/regexes/{regex_id}', response_model=StatusMessageModel)
async def regex_delete(regex_id: int):
    """Delete a regex using his id"""
    res = db.query('SELECT * FROM regexes WHERE regex_id = ?;', regex_id)
    if len(res) != 0:
        db.query('DELETE FROM regexes WHERE regex_id = ?;', regex_id)
        await firewall.get(res[0]["service_id"]).update_filters()
        await refresh_frontend()
    
    return {'status': 'ok'}

@app.post('/regexes/{regex_id}/enable', response_model=StatusMessageModel)
async def regex_enable(regex_id: int):
    """Request the enabling of a regex"""
    res = db.query('SELECT * FROM regexes WHERE regex_id = ?;', regex_id)
    if len(res) != 0:
        db.query('UPDATE regexes SET active=1 WHERE regex_id = ?;', regex_id)
        await firewall.get(res[0]["service_id"]).update_filters()
        await refresh_frontend()
    return {'status': 'ok'}

@app.post('/regexes/{regex_id}/disable', response_model=StatusMessageModel)
async def regex_disable(regex_id: int):
    """Request the deactivation of a regex"""
    res = db.query('SELECT * FROM regexes WHERE regex_id = ?;', regex_id)
    if len(res) != 0:
        db.query('UPDATE regexes SET active=0 WHERE regex_id = ?;', regex_id)
        await firewall.get(res[0]["service_id"]).update_filters()
        await refresh_frontend()
    return {'status': 'ok'}

@app.post('/regexes', response_model=StatusMessageModel)
async def add_new_regex(form: RegexAddForm):
    """Add a new regex"""
    regex_correct, message = await test_regex_validity(b64decode(form.regex))
    if not regex_correct:
        raise HTTPException(status_code=400, detail=f"Invalid regex: {message}")
    try:
        db.query("INSERT INTO regexes (service_id, regex, mode, is_case_sensitive, active ) VALUES (?, ?, ?, ?, ?);", 
                form.service_id, form.regex, form.mode, form.is_case_sensitive, True if form.active is None else form.active )
    except sqlite3.IntegrityError:
        raise HTTPException(status_code=400, detail="An identical regex already exists")

    await firewall.get(form.service_id).update_filters()
    await refresh_frontend()
    return {'status': 'ok'}

@app.post('/services', response_model=ServiceAddResponse)
async def add_new_service(form: ServiceAddForm):
    """Add a new service"""
    if form.target_type == "flow":
        if form.ip_int is None or form.port is None:
            raise HTTPException(status_code=400, detail="IP and Port are required when target type is flow")
        try:
            form.ip_int = ip_parse(form.ip_int)
        except ValueError:
            raise HTTPException(status_code=400, detail="Invalid address")
    elif form.target_type == "tls":
        if not form.tls_stream_id:
            raise HTTPException(status_code=400, detail="TLS stream ID is required when target type is tls")
            
    if form.proto not in ["tcp", "udp"]:
        raise HTTPException(status_code=400, detail="Invalid protocol")
    if form.target_type == "tls" and form.proto != "tcp":
        raise HTTPException(status_code=400, detail="TLS is only supported for TCP protocol")
            
    srv_id = None
    try:
        srv_id = gen_service_id()
        db.query("INSERT INTO services (service_id ,name, port, status, proto, ip_int, fail_open, target_type, tls_stream_id) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)",
                    srv_id, refactor_name(form.name), form.port, STATUS.STOP, form.proto, form.ip_int, form.fail_open, form.target_type, form.tls_stream_id)
    except sqlite3.IntegrityError:
        raise HTTPException(status_code=400, detail="This type of service already exists")
    await firewall.reload()
    await refresh_frontend()
    return {'status': 'ok', 'service_id': srv_id}


@app.get('/metrics', response_class = PlainTextResponse)
async def metrics():
    """Aggregate metrics"""
    stats = db.query("""
        SELECT
            s.name,
            s.status,
            r.regex,
            r.mode,
            r.is_case_sensitive,
            r.blocked_packets,
            r.active
        FROM regexes r LEFT JOIN services s ON s.service_id = r.service_id;
    """)
    metrics = []
    def sanitize(s):
        return s.replace('\\', '\\\\').replace('"', '\\"').replace('\n', '\\n')
    for stat in stats:
        props = f'service_name="{sanitize(stat["name"])}",regex="{sanitize(b64decode(stat["regex"]).decode())}",mode="{stat["mode"]}",is_case_sensitive="{stat["is_case_sensitive"]}"'
        metrics.append(f'firegex_blocked_packets{{{props}}} {stat["blocked_packets"]}')
        metrics.append(f'firegex_active{{{props}}} {int(stat["active"] and stat["status"] == "active")}')
    return "\n".join(metrics)

@app.get('/services/{service_id}/export')
async def export_service_regexes(service_id: str):
    """Export all regexes for a specific service as JSON"""
    if len(db.query('SELECT 1 FROM services WHERE service_id = ?;', service_id)) == 0:
        raise HTTPException(status_code=404, detail="Service not found")
    
    regexes = db.query("SELECT regex, mode, is_case_sensitive, active FROM regexes WHERE service_id = ?;", service_id)
    return {"regexes": regexes}

@app.post('/services/{service_id}/import', response_model=StatusMessageModel)
async def import_service_regexes(service_id: str, data: dict):
    """Import regexes for a specific service from JSON"""
    if len(db.query('SELECT 1 FROM services WHERE service_id = ?;', service_id)) == 0:
        raise HTTPException(status_code=404, detail="Service not found")
        
    regexes = data.get("regexes", [])
    if not isinstance(regexes, list):
        raise HTTPException(status_code=400, detail="Invalid format")
        
    for r in regexes:
        regex = r.get("regex")
        mode = r.get("mode")
        is_case_sensitive = r.get("is_case_sensitive", False)
        active = r.get("active", True)
        
        if not regex or not mode:
            continue
            
        try:
            db.query('INSERT INTO regexes (regex, mode, service_id, is_case_sensitive, active) VALUES (?, ?, ?, ?, ?);', 
                     regex, mode, service_id, int(is_case_sensitive), int(active))
            
            # Since we inserted it directly, we need to notify the firewall to update its rules if it's active
            if active and firewall.get(service_id).status == STATUS.ACTIVE:
                # Easiest way to sync is to restart the service rule engine
                await firewall.get(service_id).next(STATUS.ACTIVE)
        except sqlite3.IntegrityError:
            # Skip duplicates
            pass
            
    # Force reload of all rules for this service
    old_status = firewall.get(service_id).status
    await firewall.remove(service_id)
    await firewall.reload()
    await firewall.get(service_id).next(old_status)
    
    await refresh_frontend()
    return {'status': 'ok'}