#!/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)…"
if ! sudo -u postgres psql -tAc "SELECT 1 FROM pg_roles WHERE rolname = '$DB_USER'" | grep -q 1; then
    sudo -u postgres psql -c "CREATE USER $DB_USER WITH PASSWORD '$DB_PASS';"
fi

if ! sudo -u postgres psql -tAc "SELECT 1 FROM pg_database WHERE datname = '$DB_NAME'" | grep -q 1; then
    sudo -u postgres psql -c "CREATE DATABASE $DB_NAME OWNER $DB_USER;"
fi

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(),
    updated_at TIMESTAMPTZ DEFAULT NOW(),
    correlation_key TEXT UNIQUE,
    packet_id TEXT,
    packet_uid TEXT,
    correlation_source VARCHAR(32),
    skb_mark BIGINT,

    -- Interface metadata
    capture_iface VARCHAR(64),
    ingress_if VARCHAR(64),
    egress_if VARCHAR(64),
    verdict VARCHAR(32),
    verdict_reason VARCHAR(128),
    verdict_confidence VARCHAR(32),
    ingress_seen_at TIMESTAMPTZ,
    egress_seen_at TIMESTAMPTZ,
    verdict_seen_at TIMESTAMPTZ,

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

    -- VLAN
    vlan_id INTEGER,

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

    -- Transport layer
    src_port INTEGER,
    dst_port INTEGER,

    -- Packet metadata
    length INTEGER,
    raw_present BOOLEAN DEFAULT FALSE,
    capture_sources TEXT[],

    -- DPI / flow enrichment metadata
    app_protocol VARCHAR(128),
    app_master_protocol VARCHAR(128),
    app_category VARCHAR(128),
    app_confidence VARCHAR(64),
    app_hostname VARCHAR(255),
    app_is_encrypted BOOLEAN,
    app_risk_score INTEGER,
    dpi_metadata JSONB,
    capture_metadata JSONB,
    telemetry_metadata JSONB,

    -- Full packet dump
    raw BYTEA
);

ALTER TABLE packets ADD COLUMN IF NOT EXISTS updated_at TIMESTAMPTZ DEFAULT NOW();
ALTER TABLE packets ADD COLUMN IF NOT EXISTS correlation_key TEXT;
ALTER TABLE packets ADD COLUMN IF NOT EXISTS packet_id TEXT;
ALTER TABLE packets ADD COLUMN IF NOT EXISTS packet_uid TEXT;
ALTER TABLE packets ADD COLUMN IF NOT EXISTS correlation_source VARCHAR(32);
ALTER TABLE packets ADD COLUMN IF NOT EXISTS skb_mark BIGINT;
ALTER TABLE packets ADD COLUMN IF NOT EXISTS capture_iface VARCHAR(64);
ALTER TABLE packets DROP COLUMN IF EXISTS iface;
ALTER TABLE packets ADD COLUMN IF NOT EXISTS app_protocol VARCHAR(128);
ALTER TABLE packets ADD COLUMN IF NOT EXISTS app_master_protocol VARCHAR(128);
ALTER TABLE packets ADD COLUMN IF NOT EXISTS app_category VARCHAR(128);
ALTER TABLE packets ADD COLUMN IF NOT EXISTS app_confidence VARCHAR(64);
ALTER TABLE packets ADD COLUMN IF NOT EXISTS app_hostname VARCHAR(255);
ALTER TABLE packets ADD COLUMN IF NOT EXISTS app_is_encrypted BOOLEAN;
ALTER TABLE packets ADD COLUMN IF NOT EXISTS app_risk_score INTEGER;
ALTER TABLE packets ADD COLUMN IF NOT EXISTS dpi_metadata JSONB;
ALTER TABLE packets ALTER COLUMN dpi_metadata TYPE JSONB USING
    CASE
        WHEN dpi_metadata IS NULL THEN NULL
        WHEN pg_typeof(dpi_metadata)::text = 'jsonb' THEN dpi_metadata
        ELSE dpi_metadata::jsonb
    END;
ALTER TABLE packets ADD COLUMN IF NOT EXISTS capture_metadata JSONB;
ALTER TABLE packets ALTER COLUMN capture_metadata TYPE JSONB USING
    CASE
        WHEN capture_metadata IS NULL THEN NULL
        WHEN pg_typeof(capture_metadata)::text = 'jsonb' THEN capture_metadata
        ELSE capture_metadata::jsonb
    END;
ALTER TABLE packets ADD COLUMN IF NOT EXISTS eth_type_raw INTEGER;
ALTER TABLE packets ADD COLUMN IF NOT EXISTS ip_proto_raw INTEGER;
ALTER TABLE packets ADD COLUMN IF NOT EXISTS raw_present BOOLEAN DEFAULT FALSE;
ALTER TABLE packets ADD COLUMN IF NOT EXISTS capture_sources TEXT[];
ALTER TABLE packets ADD COLUMN IF NOT EXISTS ingress_if VARCHAR(64);
ALTER TABLE packets ADD COLUMN IF NOT EXISTS egress_if VARCHAR(64);
ALTER TABLE packets DROP COLUMN IF EXISTS observed_ifaces;
ALTER TABLE packets ADD COLUMN IF NOT EXISTS verdict VARCHAR(32);
ALTER TABLE packets ADD COLUMN IF NOT EXISTS verdict_reason VARCHAR(128);
ALTER TABLE packets ADD COLUMN IF NOT EXISTS verdict_confidence VARCHAR(32);
ALTER TABLE packets ADD COLUMN IF NOT EXISTS ingress_seen_at TIMESTAMPTZ;
ALTER TABLE packets ADD COLUMN IF NOT EXISTS egress_seen_at TIMESTAMPTZ;
ALTER TABLE packets ADD COLUMN IF NOT EXISTS verdict_seen_at TIMESTAMPTZ;
ALTER TABLE packets ADD COLUMN IF NOT EXISTS telemetry_metadata JSONB;
ALTER TABLE packets ALTER COLUMN telemetry_metadata TYPE JSONB USING
    CASE
        WHEN telemetry_metadata IS NULL THEN NULL
        WHEN pg_typeof(telemetry_metadata)::text = 'jsonb' THEN telemetry_metadata
        ELSE telemetry_metadata::jsonb
    END;
ALTER TABLE packets DROP COLUMN IF EXISTS direction;
UPDATE packets
SET raw_present = (raw IS NOT NULL)
WHERE raw_present IS DISTINCT FROM (raw IS NOT NULL);
UPDATE packets
SET correlation_key = CASE
    WHEN packet_id IS NOT NULL AND packet_id <> '' THEN 'pid:' || packet_id
    WHEN packet_uid IS NOT NULL AND packet_uid <> '' THEN 'uid:' || packet_uid
    ELSE 'row:' || id::text
END
WHERE correlation_key IS NULL OR correlation_key = '';
UPDATE packets
SET correlation_source = CASE
    WHEN packet_id IS NOT NULL AND packet_id <> '' THEN 'kernel_mark'
    ELSE 'legacy_hash'
END
WHERE correlation_source IS NULL OR correlation_source = '';
UPDATE packets
SET capture_sources = ARRAY_REMOVE(ARRAY[
    CASE WHEN raw IS NOT NULL THEN 'af_packet' END,
    CASE WHEN telemetry_metadata IS NOT NULL THEN 'telemetry' END
], NULL)
WHERE capture_sources IS NULL OR array_length(capture_sources, 1) IS NULL;
UPDATE packets
SET updated_at = COALESCE(verdict_seen_at, egress_seen_at, ingress_seen_at, timestamp, NOW())
WHERE updated_at IS NULL;
CREATE UNIQUE INDEX IF NOT EXISTS idx_packets_correlation_key ON packets(correlation_key);
CREATE UNIQUE INDEX IF NOT EXISTS idx_packets_packet_uid ON packets(packet_uid);
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;

-- Change ownership to ensure the user can manage the table
ALTER TABLE packets OWNER TO $DB_USER;
ALTER SEQUENCE packets_id_seq OWNER 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;

-- Create indexes for faster queries
CREATE INDEX IF NOT EXISTS idx_packets_timestamp ON packets(timestamp DESC);
CREATE INDEX IF NOT EXISTS idx_packets_updated_at ON packets(updated_at DESC);
CREATE INDEX IF NOT EXISTS idx_packets_packet_id ON packets(packet_id);
CREATE INDEX IF NOT EXISTS idx_packets_src_ip ON packets(src_ip);
CREATE INDEX IF NOT EXISTS idx_packets_app_protocol ON packets(app_protocol);
CREATE INDEX IF NOT EXISTS idx_packets_app_hostname ON packets(app_hostname);
CREATE INDEX IF NOT EXISTS idx_packets_verdict ON packets(verdict);
EOF

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