#!/bin/bash
set -e

DB_NAME="mitm_db"
DB_USER="mitm_user"
DB_PASS="mitm_password"
LAN_SUBNET="192.168.178.0/24"

echo "[1] Installing PostgreSQL…"
sudo apt update -y
sudo apt install -y postgresql postgresql-contrib

echo "[2] Configuring PostgreSQL to listen on all addresses…"
PG_CONF="/etc/postgresql/$(ls /etc/postgresql)/main/postgresql.conf"
sudo sed -i "s/^#listen_addresses =.*/listen_addresses = '*'/" "$PG_CONF"

echo "[3] Updating pg_hba.conf for LAN + localhost access…"
PG_HBA="/etc/postgresql/$(ls /etc/postgresql)/main/pg_hba.conf"

# Add localhost
if ! grep -Eq "^[ ]*host[ ]+all[ ]+all[ ]+127.0.0.1/32" "$PG_HBA"; then
    echo "host    all     all     127.0.0.1/32        md5" | sudo tee -a "$PG_HBA"
fi

# Add LAN
if ! grep -q "$LAN_SUBNET" "$PG_HBA"; then
    echo "host    all     all     $LAN_SUBNET        md5" | sudo tee -a "$PG_HBA"
fi

echo "[4] Restarting PostgreSQL…"
sudo systemctl restart postgresql

echo "[5] Creating database + user (idempotent)…"
sudo -u postgres psql <<EOF
DO \$\$
BEGIN
    IF NOT EXISTS (SELECT FROM pg_roles WHERE rolname = '$DB_USER') THEN
        CREATE USER $DB_USER WITH PASSWORD '$DB_PASS';
    END IF;

    IF NOT EXISTS (SELECT FROM pg_database WHERE datname = '$DB_NAME') THEN
        CREATE DATABASE $DB_NAME OWNER $DB_USER;
    END IF;
END
\$\$;
EOF

echo "[6] Creating packets table (idempotent)…"
sudo -u postgres psql -d "$DB_NAME" <<EOF
CREATE TABLE IF NOT EXISTS packets (
    id BIGSERIAL PRIMARY KEY,
    timestamp TIMESTAMPTZ DEFAULT NOW(),

    -- Interface metadata
    iface VARCHAR(64),
    direction VARCHAR(64),

    -- Ethernet
    src_mac MACADDR,
    dst_mac MACADDR,
    eth_type VARCHAR(64),

    -- VLAN
    vlan_id INTEGER,

    -- IP layer
    src_ip INET,
    dst_ip INET,
    ip_proto VARCHAR(64),

    -- Transport layer
    src_port INTEGER,
    dst_port INTEGER,

    -- Packet metadata
    length INTEGER,

    -- Full packet dump
    raw BYTEA
);
EOF


echo "[7] Grant privileges to user…"
sudo -u postgres psql -d $DB_NAME <<EOF
-- Grant all table privileges
GRANT ALL PRIVILEGES ON ALL TABLES IN SCHEMA public TO $DB_USER;

-- Grant sequence permissions (needed for SERIAL / BIGSERIAL)
GRANT USAGE, SELECT, UPDATE ON ALL SEQUENCES IN SCHEMA public TO $DB_USER;

-- Ensure future tables & sequences also get permissions
ALTER DEFAULT PRIVILEGES IN SCHEMA public GRANT ALL PRIVILEGES ON TABLES TO $DB_USER;
ALTER DEFAULT PRIVILEGES IN SCHEMA public GRANT USAGE, SELECT, UPDATE ON SEQUENCES TO $DB_USER;
EOF

echo "Done. PostgreSQL is ready for LAN + localhost connections."