#!/bin/bash
set -e

DB_NAME="mitm_traffic"
DB_USER="mitm_user"
DB_PASS="mitm_password"

echo "=== MITM Traffic DB Setup ==="

# ---------------------------
# Install PostgreSQL (idempotent)
# ---------------------------
if ! command -v psql >/dev/null 2>&1; then
    echo "[+] Installing PostgreSQL..."
    sudo apt update
    sudo apt install -y postgresql postgresql-contrib
else
    echo "[+] PostgreSQL already installed"
fi

# Ensure PostgreSQL service is running
sudo systemctl enable postgresql
sudo systemctl start postgresql

# ---------------------------
# Create DB user (idempotent)
# ---------------------------
USER_EXISTS=$(sudo -u postgres psql -tAc "SELECT 1 FROM pg_roles WHERE rolname='${DB_USER}'")

if [ "$USER_EXISTS" = "1" ]; then
    echo "[+] DB user already exists"
else
    echo "[+] Creating DB user..."
    sudo -u postgres psql -c "CREATE USER ${DB_USER} WITH PASSWORD '${DB_PASS}';"
fi

# ---------------------------
# Create database (idempotent)
# ---------------------------
DB_EXISTS=$(sudo -u postgres psql -tAc "SELECT 1 FROM pg_database WHERE datname='${DB_NAME}'")

if [ "$DB_EXISTS" = "1" ]; then
    echo "[+] Database already exists"
else
    echo "[+] Creating database..."
    sudo -u postgres psql -c "CREATE DATABASE ${DB_NAME} OWNER ${DB_USER};"
fi

# ---------------------------
# Create tables (idempotent)
# ---------------------------

echo "[+] Creating tables…"

sudo -u postgres psql -d "$DB_NAME" << 'EOF'

-- Ingress und Egress Pakete werden getrennt gespeichert
CREATE TABLE IF NOT EXISTS ingress_packets (
    id BIGSERIAL PRIMARY KEY,
    timestamp TIMESTAMPTZ DEFAULT NOW(),
    iface VARCHAR(32),
    src_mac VARCHAR(32),
    dst_mac VARCHAR(32),
    eth_type VARCHAR(8),
    src_ip VARCHAR(64),
    dst_ip VARCHAR(64),
    src_port INTEGER,
    dst_port INTEGER,
    protocol VARCHAR(16),
    packet_length INTEGER,
    raw BYTEA
);

CREATE TABLE IF NOT EXISTS egress_packets (
    id BIGSERIAL PRIMARY KEY,
    timestamp TIMESTAMPTZ DEFAULT NOW(),
    iface VARCHAR(32),
    src_mac VARCHAR(32),
    dst_mac VARCHAR(32),
    eth_type VARCHAR(8),
    src_ip VARCHAR(64),
    dst_ip VARCHAR(64),
    src_port INTEGER,
    dst_port INTEGER,
    protocol VARCHAR(16),
    packet_length INTEGER,
    raw BYTEA
);

-- eBPF / nftables / ebtables decisions
CREATE TABLE IF NOT EXISTS packet_verdicts (
    id BIGSERIAL PRIMARY KEY,
    timestamp TIMESTAMPTZ DEFAULT NOW(),
    packet_hash VARCHAR(128),  -- Referenz zu ingress/egress
    table VARCHAR(64),
    chain VARCHAR(64),
    rule_id INTEGER,
    verdict VARCHAR(32),        -- ACCEPT / DROP / STOLEN / QUEUE / ...
    metadata JSONB
);

-- Optionaler Index für Performance
CREATE INDEX IF NOT EXISTS idx_verdict_packet_hash ON packet_verdicts(packet_hash);
CREATE INDEX IF NOT EXISTS idx_ingress_timestamp ON ingress_packets(timestamp);
CREATE INDEX IF NOT EXISTS idx_egress_timestamp ON egress_packets(timestamp);

EOF

echo "[+] All done! The database is ready."