giglez/scripts/import_tembed_signatures.py

#!/usr/bin/env python3
"""
Import T-Embed RF captures as signature database entries

This script:
1. Scans T-Embed .sub files
2. Extracts RF signal patterns
3. Creates device and signature records
4. Populates database for matching

Based on real wardriving captures from T-Embed device
"""

import sys
import json
from pathlib import Path
from datetime import datetime
from typing import List, Dict, Any, Optional

# Add project root to path
sys.path.insert(0, str(Path(__file__).parent.parent))

from src.parser.sub_parser import SubFileParser
from src.database.models import Device, Signature, FlipperSignature, Base
from src.database.connection import get_engine, get_session
from sqlalchemy.orm import Session


class TembedSignatureImporter:
    """Import T-Embed RF captures as device signatures"""

    def __init__(self, session: Session):
        self.session = session
        self.parser = SubFileParser()
        self.stats = {
            'files_found': 0,
            'files_parsed': 0,
            'files_skipped': 0,
            'devices_created': 0,
            'signatures_created': 0,
            'errors': []
        }

    def import_directory(self, directory: Path) -> Dict[str, Any]:
        """
        Import all .sub files from directory

        Args:
            directory: Path to directory containing .sub files

        Returns:
            Statistics dictionary
        """
        print(f"Scanning {directory} for .sub files...")

        # Find all .sub files
        sub_files = list(directory.glob('*.sub'))
        self.stats['files_found'] = len(sub_files)

        print(f"Found {len(sub_files)} .sub files\n")

        for sub_file in sorted(sub_files):
            print(f"Processing: {sub_file.name}")
            try:
                self._import_file(sub_file)
            except Exception as e:
                error_msg = f"Error processing {sub_file.name}: {e}"
                print(f"  ❌ {error_msg}")
                self.stats['errors'].append(error_msg)

        # Commit all changes
        try:
            self.session.commit()
            print("\n✅ Database changes committed")
        except Exception as e:
            self.session.rollback()
            print(f"\n❌ Failed to commit: {e}")
            self.stats['errors'].append(f"Commit failed: {e}")

        return self.stats

    def _import_file(self, file_path: Path):
        """Import a single .sub file"""

        # Parse file
        try:
            metadata = self.parser.parse(str(file_path))
        except Exception as e:
            self.stats['files_skipped'] += 1
            raise ValueError(f"Parse failed: {e}")

        # Skip if empty (frequency = 0, no data)
        if metadata.frequency == 0 or (metadata.file_format == 'RAW' and not metadata.raw_data):
            print(f"  ⏭️  Skipped: Empty capture")
            self.stats['files_skipped'] += 1
            return

        self.stats['files_parsed'] += 1

        # Load GPS data if available
        gps_data = self._load_gps_data(file_path)

        # Create device record
        device = self._create_device(metadata, file_path, gps_data)

        # Create signature record
        signature = self._create_signature(metadata, device)

        # Create Flipper signature record (for compatibility)
        flipper_sig = self._create_flipper_signature(metadata, device, file_path)

        print(f"  ✅ Device: {device.device_name}")
        print(f"     Frequency: {metadata.frequency/1e6:.2f} MHz")
        if metadata.raw_data:
            print(f"     RAW samples: {len(metadata.raw_data)}")
        if gps_data:
            print(f"     GPS: {gps_data['data']['latitude']:.4f}, {gps_data['data']['longitude']:.4f}")

    def _load_gps_data(self, sub_file: Path) -> Optional[Dict]:
        """Load GPS coordinates for a .sub file if available"""

        # Look for matching GPS JSON files in same directory
        # Pattern: gps_coordinates_YYYYMMDD_HHMMSS.json
        gps_files = sorted(sub_file.parent.glob('gps_coordinates_*.json'))

        if not gps_files:
            return None

        # Use the most recent GPS file (simple heuristic)
        gps_file = gps_files[-1]

        try:
            with open(gps_file, 'r') as f:
                return json.load(f)
        except Exception as e:
            print(f"  ⚠️  Could not load GPS data: {e}")
            return None

    def _create_device(self, metadata, file_path: Path, gps_data: Optional[Dict]) -> Device:
        """Create a Device record from metadata"""

        # Generate device name from file and frequency
        device_name = self._generate_device_name(file_path, metadata)

        # Determine device type from frequency
        device_type = self._guess_device_type(metadata.frequency)

        # Check if device already exists
        existing = self.session.query(Device).filter(
            Device.device_name == device_name
        ).first()

        if existing:
            print(f"  ℹ️  Device already exists: {device_name}")
            return existing

        # Create new device
        device = Device(
            device_name=device_name,
            manufacturer='Unknown',
            model='T-Embed Capture',
            device_type=device_type,
            typical_frequency=metadata.frequency,
            protocol=metadata.protocol if metadata.protocol else 'RAW',
            description=f"Captured from T-Embed RF device at {file_path.name}",
            first_seen=datetime.utcnow(),
            is_verified=False
        )

        self.session.add(device)
        self.session.flush()  # Get device.id

        self.stats['devices_created'] += 1

        return device

    def _create_signature(self, metadata, device: Device) -> Signature:
        """Create a Signature record for pattern matching"""

        # Check if signature already exists
        existing = self.session.query(Signature).filter(
            Signature.device_id == device.id,
            Signature.frequency == metadata.frequency
        ).first()

        if existing:
            return existing

        # Extract timing patterns for RAW format
        timing_min, timing_max = None, None
        if metadata.raw_data and len(metadata.raw_data) > 0:
            # Use absolute values for timing
            abs_timings = [abs(t) for t in metadata.raw_data]
            timing_min = min(abs_timings)
            timing_max = max(abs_timings)

        signature = Signature(
            device_id=device.id,
            protocol=metadata.protocol if metadata.protocol else 'RAW',
            frequency=metadata.frequency,
            modulation=metadata.modulation,
            bit_pattern=None,  # Not available for RAW
            bit_mask=None,
            timing_min=timing_min,
            timing_max=timing_max,
            raw_pattern=','.join(map(str, metadata.raw_data)) if metadata.raw_data else None,
            confidence_threshold=0.7,  # Default threshold
            source='tembed_wardriving',
            created_at=datetime.utcnow()
        )

        self.session.add(signature)
        self.stats['signatures_created'] += 1

        return signature

    def _create_flipper_signature(self, metadata, device: Device, file_path: Path) -> FlipperSignature:
        """Create FlipperSignature record for compatibility"""

        # Check if exists
        existing = self.session.query(FlipperSignature).filter(
            FlipperSignature.device_id == device.id
        ).first()

        if existing:
            return existing

        flipper_sig = FlipperSignature(
            device_id=device.id,
            frequency=metadata.frequency,
            preset=metadata.preset if metadata.preset else '0',
            protocol=metadata.protocol if metadata.protocol else 'RAW',
            bit=metadata.bit_length,
            key=metadata.key_data,
            te=metadata.timing_element,
            raw_data=','.join(map(str, metadata.raw_data)) if metadata.raw_data else None,
            source_file=file_path.name,
            imported_at=datetime.utcnow()
        )

        self.session.add(flipper_sig)

        return flipper_sig

    def _generate_device_name(self, file_path: Path, metadata) -> str:
        """Generate a unique device name"""
        freq_mhz = metadata.frequency / 1e6
        return f"{file_path.stem}_{freq_mhz:.0f}MHz"

    def _guess_device_type(self, frequency: int) -> str:
        """Guess device type from frequency"""
        freq_mhz = frequency / 1e6

        if 300 <= freq_mhz <= 350:
            return 'garage_door'
        elif 400 <= freq_mhz <= 440:
            return 'remote_control'
        elif 860 <= freq_mhz <= 870:
            return 'sensor'
        elif 900 <= freq_mhz <= 930:
            return 'ism_device'  # 915 MHz ISM band
        else:
            return 'unknown'

    def print_summary(self):
        """Print import summary"""
        print("\n" + "="*60)
        print("IMPORT SUMMARY")
        print("="*60)
        print(f"Files found:       {self.stats['files_found']}")
        print(f"Files parsed:      {self.stats['files_parsed']}")
        print(f"Files skipped:     {self.stats['files_skipped']}")
        print(f"Devices created:   {self.stats['devices_created']}")
        print(f"Signatures created: {self.stats['signatures_created']}")

        if self.stats['errors']:
            print(f"\nErrors: {len(self.stats['errors'])}")
            for error in self.stats['errors']:
                print(f"  - {error}")
        else:
            print("\n✅ No errors")

        print("="*60)


def main():
    """Main entry point"""

    print("="*60)
    print("T-Embed RF Signature Importer")
    print("="*60)
    print()

    # Get database session
    try:
        engine = get_engine()
        session = get_session()
        print("✅ Database connected")
    except Exception as e:
        print(f"❌ Database connection failed: {e}")
        print("\nMake sure PostgreSQL is running and configured correctly")
        return 1

    # Create tables if needed
    try:
        Base.metadata.create_all(engine)
        print("✅ Database tables ready\n")
    except Exception as e:
        print(f"⚠️  Could not create tables: {e}\n")

    # Find T-Embed directory
    tembed_dir = Path(__file__).parent.parent / 'signatures' / 't-embed-rf'

    if not tembed_dir.exists():
        print(f"❌ Directory not found: {tembed_dir}")
        return 1

    # Import signatures
    importer = TembedSignatureImporter(session)
    stats = importer.import_directory(tembed_dir)
    importer.print_summary()

    session.close()

    return 0 if not stats['errors'] else 1


if __name__ == '__main__':
    sys.exit(main())