fsk_decoder.py

"""
fsk_decoder.py — FSK Signal Demodulator + RNA-to-ASCII Decoder
Project Helix CTF · TCM Security
Author: Dibyanshu Sekhar

Decodes a 4-frequency FSK signal encoding RNA bases into ASCII via quaternary math.

Signal spec (from Owens' LOG_088):
  - Four frequencies: 1000, 2000, 3000, 4000 Hz → bases A, C, G, U
  - Base duration: ~0.85s · Silence gap between codons: ~0.75s
  - Each codon: 4 bases → ASCII = b0×64 + b1×16 + b2×4 + b3
  - Message repeats every ~78s with | (CUUA) as delimiters

Usage:
    python fsk_decoder.py --input /tmp/signal_long.raw [--verbose]

Dependencies:
    pip install numpy
"""

import numpy as np
import argparse
from collections import Counter


SAMPLE_RATE = 44100
FFT_SIZE    = 2048
HOP_SIZE    = FFT_SIZE // 2

# FSK frequency centers
FSK_CENTERS = [1000, 2000, 3000, 4000]
FREQ_TO_BASE = {1000: 'A', 2000: 'C', 3000: 'G', 4000: 'U'}
BASE_TO_VAL  = {'A': 0, 'C': 1, 'G': 2, 'U': 3}

# Timing thresholds (seconds)
MIN_SEGMENT_DURATION = 0.15   # ignore noise blips shorter than this
CODON_SIZE           = 4      # bases per codon


def load_raw_pcm(path: str) -> np.ndarray:
    """Load raw Int16 PCM mono audio."""
    raw = np.frombuffer(open(path, 'rb').read(), dtype=np.int16)
    return raw.astype(np.float32)


def detect_frames(audio: np.ndarray, verbose: bool = False) -> list:
    """
    Slide an FFT window over audio and classify each frame as a base or silence.
    Returns a list of (time_seconds, base_or_None) per frame.
    """
    freqs = np.fft.rfftfreq(FFT_SIZE, 1 / SAMPLE_RATE)
    window = np.hanning(FFT_SIZE)
    results = []

    n_frames = (len(audio) - FFT_SIZE) // HOP_SIZE

    for i in range(n_frames):
        start = i * HOP_SIZE
        frame = audio[start: start + FFT_SIZE]
        spec  = np.abs(np.fft.rfft(frame * window))

        # Measure energy in a ±200 Hz band around each center frequency
        energies = {}
        for fc in FSK_CENTERS:
            mask = (freqs >= fc - 200) & (freqs <= fc + 200)
            energies[fc] = spec[mask].max()

        total = sum(energies.values())
        best  = max(energies, key=energies.get)

        # Silence heuristic: dominant energy too weak or not clearly dominant
        if energies[best] < 3000 or energies[best] < total * 0.25:
            results.append((start / SAMPLE_RATE, None))
        else:
            results.append((start / SAMPLE_RATE, best))

    if verbose:
        print(f"[*] Analyzed {n_frames:,} frames")

    return results


def frames_to_segments(frames: list) -> list:
    """
    Merge consecutive same-frequency frames into segments.
    Returns list of (start_time, end_time, frequency_or_None).
    """
    if not frames:
        return []

    segments = []
    cur_time, cur_freq = frames[0]
    seg_start = cur_time

    for time_s, freq in frames[1:]:
        if freq != cur_freq:
            segments.append((seg_start, time_s, cur_freq))
            seg_start = time_s
            cur_freq = freq

    if cur_freq is not None:
        segments.append((seg_start, frames[-1][0], cur_freq))

    # Filter out noise (too short)
    segments = [
        s for s in segments
        if (s[1] - s[0]) >= MIN_SEGMENT_DURATION
    ]

    return segments


def segments_to_codons(segments: list, verbose: bool = False) -> list:
    """
    Group non-silence segments between silence gaps into codons.
    Returns list of codon strings like 'CUUA'.
    """
    codons = []
    current_codon_bases = []

    for start, end, freq in segments:
        duration = end - start

        if freq is None:
            # Silence = codon boundary
            if current_codon_bases:
                codon = ''.join(current_codon_bases)
                codons.append(codon)
                if verbose:
                    print(f"  Codon: {codon}")
                current_codon_bases = []
        else:
            base = FREQ_TO_BASE[freq]
            # Estimate how many base-lengths fit in this segment (~0.85s each)
            n_bases = max(1, round(duration / 0.85))
            current_codon_bases.extend([base] * n_bases)

            # Flush if we've hit codon size
            while len(current_codon_bases) >= CODON_SIZE:
                codon = ''.join(current_codon_bases[:CODON_SIZE])
                codons.append(codon)
                if verbose:
                    print(f"  Codon: {codon}")
                current_codon_bases = current_codon_bases[CODON_SIZE:]

    # Flush any trailing bases
    if current_codon_bases:
        codon = ''.join(current_codon_bases)
        codons.append(codon)

    return codons


def decode_codon(codon: str) -> str | None:
    """Decode a 4-base RNA codon to an ASCII character."""
    if len(codon) != CODON_SIZE:
        return None
    try:
        vals = [BASE_TO_VAL[b] for b in codon]
        ascii_val = vals[0] * 64 + vals[1] * 16 + vals[2] * 4 + vals[3]
        if 32 <= ascii_val <= 126:
            return chr(ascii_val)
    except KeyError:
        pass
    return None


def extract_clean_message(decoded: str) -> str:
    """
    Find the cleanest complete cycle in the decoded string.
    Looks for the pattern |...|...|
    """
    if '|' not in decoded:
        return decoded

    # Find all pipe positions
    pipes = [i for i, c in enumerate(decoded) if c == '|']
    if len(pipes) < 2:
        return decoded

    # Extract between first two pipes for one clean cycle
    start = pipes[0]
    end   = pipes[1] + 1
    return decoded[start:end]


def main():
    parser = argparse.ArgumentParser(description='FSK Decoder — Project Helix')
    parser.add_argument('--input', required=True, help='Path to raw PCM file')
    parser.add_argument('--verbose', action='store_true', help='Show per-codon decoding')
    args = parser.parse_args()

    print(f"[*] Loading: {args.input}")
    audio = load_raw_pcm(args.input)
    print(f"[*] Duration: {len(audio) / SAMPLE_RATE:.1f}s ({len(audio):,} samples)")

    print("[*] Running FFT frame analysis...")
    frames = detect_frames(audio, verbose=args.verbose)

    print("[*] Merging frames into segments...")
    segments = frames_to_segments(frames)
    non_silence = [s for s in segments if s[2] is not None]
    print(f"[*] {len(non_silence)} non-silence segments found")

    print("[*] Extracting codons...")
    codons = segments_to_codons(segments, verbose=args.verbose)
    print(f"[*] {len(codons)} codons found")

    print("\n[*] Decoding RNA → ASCII...")
    chars = []
    for i, codon in enumerate(codons):
        ch = decode_codon(codon)
        if ch:
            chars.append(ch)
            if args.verbose:
                vals = [BASE_TO_VAL.get(b, '?') for b in codon]
                val  = sum(v * (64 // (4 ** j)) for j, v in enumerate(vals)) if len(vals) == 4 else '?'
                print(f"  #{i+1:>3} {codon}  {val:>3}  '{ch}'")

    raw_decoded = ''.join(chars)
    print(f"\n[+] Raw decoded string: {raw_decoded}")

    clean = extract_clean_message(raw_decoded)
    print(f"[+] Clean message:      {clean}")
    print(f"\n[+] Decoded key: {clean}")


if __name__ == '__main__':
    main()