wiki/tasks/flightradar24/ingest/capture/main.py

"""
FR24 Capture Service
Connects to PostgreSQL, creates a capture session, and writes fake raw_packets.
In production: replace the fake packet loop with actual RTL-SDR / dump1090 input.
"""
import os
import time
import base64
import random
import logging
import signal
import sys
from datetime import datetime, timezone

import psycopg2
import psycopg2.extras

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [capture] %(levelname)s %(message)s",
    datefmt="%Y-%m-%dT%H:%M:%S",
)
log = logging.getLogger("capture")

# ── config ────────────────────────────────────────────────────────────────────
DB_DSN = (
    f"host={os.environ['POSTGRES_HOST']} "
    f"port={os.environ.get('POSTGRES_PORT', 5432)} "
    f"dbname={os.environ['POSTGRES_DB']} "
    f"user={os.environ['POSTGRES_USER']} "
    f"password={os.environ['POSTGRES_PASSWORD']}"
)
CENTER_FREQ = int(os.environ.get("RTLSDR_CENTER_FREQUENCY", 1090000000))
SAMPLE_RATE = int(os.environ.get("RTLSDR_SAMPLE_RATE", 2000000))
DEVICE_INDEX = int(os.environ.get("RTLSDR_DEVICE_INDEX", 0))
GAIN_RAW = os.environ.get("RTLSDR_GAIN", "auto")
GAIN_DB = None if GAIN_RAW == "auto" else float(GAIN_RAW)
PACKET_INTERVAL = float(os.environ.get("PACKET_INTERVAL_SECONDS", 2.0))
HEALTHCHECK_FILE = "/tmp/capture-ready"

# ── fake ADS-B payload generator ─────────────────────────────────────────────
# Real ADS-B Mode-S messages are 7 or 14 bytes.
# We generate plausible-looking random bytes tagged as DF17 (extended squitter).
_ICAO_POOL = [f"{i:06X}" for i in random.sample(range(0x400000, 0xFFFFFF), 20)]

def _fake_adsb_bytes() -> bytes:
    """14-byte fake Mode-S extended squitter (DF17)."""
    df17_first_byte = 0x8D  # downlink format 17
    icao = bytes.fromhex(random.choice(_ICAO_POOL))
    payload = bytes([random.randint(0, 255) for _ in range(7)])
    crc = bytes([random.randint(0, 255) for _ in range(3)])
    return bytes([df17_first_byte]) + icao + payload + crc

def _fake_packet_row(capture_id: str) -> dict:
    raw = _fake_adsb_bytes()
    return {
        "capture_id": capture_id,
        "observed_at": datetime.now(timezone.utc),
        "partition_date": datetime.now(timezone.utc).date(),
        "frequency_hz": CENTER_FREQ,
        "rssi_dbm": round(random.uniform(-90.0, -40.0), 3),
        "snr_db": round(random.uniform(5.0, 30.0), 3),
        "samplerate_hz": SAMPLE_RATE,
        "payload_base64": base64.b64encode(raw).decode(),
        "payload_bytes": len(raw),
        "decoded_format": "mode-s",
        "message_type": "DF17",
    }

# ── db helpers ────────────────────────────────────────────────────────────────
def wait_for_db(max_attempts: int = 30) -> psycopg2.extensions.connection:
    for attempt in range(1, max_attempts + 1):
        try:
            conn = psycopg2.connect(DB_DSN)
            log.info("PostgreSQL connected (attempt %d)", attempt)
            return conn
        except psycopg2.OperationalError as e:
            log.warning("DB not ready (%d/%d): %s", attempt, max_attempts, e)
            time.sleep(2)
    log.error("Could not connect to PostgreSQL after %d attempts", max_attempts)
    sys.exit(1)

def create_capture_session(conn) -> str:
    with conn.cursor() as cur:
        cur.execute(
            """
            INSERT INTO fr24.captures
              (started_at, source, device_index, center_frequency_hz,
               sample_rate_hz, gain_db, status, notes)
            VALUES (%s, %s, %s, %s, %s, %s, 'active', %s)
            RETURNING capture_id
            """,
            (
                datetime.now(timezone.utc),
                "rtl-sdr",
                DEVICE_INDEX,
                CENTER_FREQ,
                SAMPLE_RATE,
                GAIN_DB,
                "fake-data mode (step-1 scaffold)",
            ),
        )
        capture_id = str(cur.fetchone()[0])
    conn.commit()
    log.info("Capture session created: %s", capture_id)
    return capture_id

def insert_packet(conn, row: dict):
    with conn.cursor() as cur:
        cur.execute(
            """
            INSERT INTO fr24.raw_packets
              (capture_id, observed_at, partition_date, frequency_hz,
               rssi_dbm, snr_db, samplerate_hz, payload_base64,
               payload_bytes, decoded_format, message_type)
            VALUES
              (%(capture_id)s, %(observed_at)s, %(partition_date)s, %(frequency_hz)s,
               %(rssi_dbm)s, %(snr_db)s, %(samplerate_hz)s, %(payload_base64)s,
               %(payload_bytes)s, %(decoded_format)s, %(message_type)s)
            """,
            row,
        )
    conn.commit()

def close_capture_session(conn, capture_id: str):
    with conn.cursor() as cur:
        cur.execute(
            "UPDATE fr24.captures SET ended_at=%s, status='stopped', updated_at=now() WHERE capture_id=%s",
            (datetime.now(timezone.utc), capture_id),
        )
    conn.commit()
    log.info("Capture session closed: %s", capture_id)

# ── main ──────────────────────────────────────────────────────────────────────
def main():
    conn = wait_for_db()
    capture_id = create_capture_session(conn)

    # signal healthcheck
    open(HEALTHCHECK_FILE, "w").close()
    log.info("Healthcheck file written: %s", HEALTHCHECK_FILE)

    shutdown = {"flag": False}

    def _handle_signal(sig, frame):
        log.info("Signal %s received, shutting down", sig)
        shutdown["flag"] = True

    signal.signal(signal.SIGTERM, _handle_signal)
    signal.signal(signal.SIGINT, _handle_signal)

    packet_count = 0
    log.info("Starting fake packet loop (interval=%.1fs)", PACKET_INTERVAL)

    while not shutdown["flag"]:
        try:
            row = _fake_packet_row(capture_id)
            insert_packet(conn, row)
            packet_count += 1
            if packet_count % 10 == 0:
                log.info("Packets written: %d", packet_count)
        except Exception as e:
            log.error("Packet insert failed: %s", e)
            try:
                conn.rollback()
            except Exception:
                pass

        time.sleep(PACKET_INTERVAL)

    close_capture_session(conn, capture_id)
    conn.close()
    log.info("Capture service stopped. Total packets: %d", packet_count)

if __name__ == "__main__":
    main()