#!/usr/bin/env bash
set -e

DB_NAME="mitm_db"
DB_USER="mitm_user"
DB_PASS="mitm_pass"

echo "[+] Installing PostgreSQL..."
sudo apt-get update -y
sudo apt-get install -y postgresql postgresql-contrib

echo "[+] Ensuring PostgreSQL is running..."
sudo systemctl enable postgresql
sudo systemctl start postgresql

echo "[+] Creating user and database (if missing)..."
sudo -u postgres psql <<EOF
DO
\$do\$
BEGIN
   IF NOT EXISTS (SELECT FROM pg_catalog.pg_roles WHERE rolname = '${DB_USER}') THEN
      CREATE ROLE ${DB_USER} LOGIN PASSWORD '${DB_PASS}';
   END IF;
END
\$do\$;

CREATE DATABASE ${DB_NAME} OWNER ${DB_USER}
    TEMPLATE template0
    ENCODING 'UTF8'
    LC_COLLATE='C'
    LC_CTYPE='C'
    CONNECTION LIMIT -1;

-- ignore errors if DB already exists
ALTER DATABASE ${DB_NAME} OWNER TO ${DB_USER};
EOF

echo "[+] Creating tables (idempotent)..."
sudo -u postgres psql -d ${DB_NAME} <<EOF

CREATE TABLE IF NOT EXISTS packets (
    id BIGSERIAL PRIMARY KEY,
    timestamp TIMESTAMPTZ DEFAULT NOW(),

    direction VARCHAR(16) NOT NULL,
    iface VARCHAR(32) NOT NULL,

    src_mac VARCHAR(17),
    dst_mac VARCHAR(17),

    eth_type VARCHAR(8),

    src_ip VARCHAR(64),
    dst_ip VARCHAR(64),
    ip_proto VARCHAR(16),

    src_port INT,
    dst_port INT,

    packet_len INT,
    packet_hex TEXT
);

CREATE TABLE IF NOT EXISTS nft_events (
    id BIGSERIAL PRIMARY KEY,
    timestamp TIMESTAMPTZ DEFAULT NOW(),
    packet_id BIGINT REFERENCES packets(id) ON DELETE CASCADE,

    hook VARCHAR(64),
    nft_table VARCHAR(64),
    nft_chain VARCHAR(64),
    nft_rule VARCHAR(256),
    verdict VARCHAR(32)
);
EOF

echo "[+] DONE!"
echo "Database: ${DB_NAME}"
echo "User:     ${DB_USER}"
echo "Password: ${DB_PASS}"