giglez/scripts/create_test_dataset.py

#!/usr/bin/env python3
"""
Create Test Dataset with GPS Coordinates

Adds GPS coordinates to existing .sub files from Flipper Zero for testing
"""

import sys
import json
import shutil
from pathlib import Path
from datetime import datetime, timezone

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

# Test GPS coordinates across different US cities
TEST_LOCATIONS = {
    "los_angeles": {
        "name": "Los Angeles, CA",
        "coords": [(34.0522, -118.2437), (34.0736, -118.3596), (34.0478, -118.2348)]
    },
    "san_francisco": {
        "name": "San Francisco, CA",
        "coords": [(37.7749, -122.4194), (37.8044, -122.2712), (37.7558, -122.4449)]
    },
    "new_york": {
        "name": "New York, NY",
        "coords": [(40.7128, -74.0060), (40.7580, -73.9855), (40.6782, -73.9442)]
    },
    "chicago": {
        "name": "Chicago, IL",
        "coords": [(41.8781, -87.6298), (41.8919, -87.6051), (41.8369, -87.6847)]
    },
    "austin": {
        "name": "Austin, TX",
        "coords": [(30.2672, -97.7431), (30.3072, -97.7559), (30.2500, -97.7500)]
    },
}


def create_gps_tagged_captures():
    """Copy Flipper Zero captures and add GPS tags via filenames and JSON"""

    # Source directory
    flipper_dir = Path("signatures/flipperzero-firmware/applications/debug/unit_tests/resources/unit_tests/subghz")

    # Output directory
    output_dir = Path("signatures/test_dataset_gps")
    output_dir.mkdir(exist_ok=True)

    print("=" * 80)
    print("Creating Test Dataset with GPS Coordinates")
    print("=" * 80)

    if not flipper_dir.exists():
        print(f"❌ Flipper directory not found: {flipper_dir}")
        return

    # Get all .sub files
    sub_files = list(flipper_dir.glob("*.sub"))

    print(f"\n📁 Found {len(sub_files)} .sub files")
    print(f"📍 Creating GPS-tagged versions across {len(TEST_LOCATIONS)} cities")

    created_count = 0
    location_idx = 0

    # Create captures for each location
    for city_key, city_data in TEST_LOCATIONS.items():
        city_name = city_data["name"]
        coords_list = city_data["coords"]

        print(f"\n🌎 {city_name}")

        for i, (lat, lon) in enumerate(coords_list):
            if location_idx >= len(sub_files):
                break

            source_file = sub_files[location_idx]

            # Create GPS-tagged filename
            lat_str = f"{abs(lat):.4f}{'N' if lat >= 0 else 'S'}"
            lon_str = f"{abs(lon):.4f}{'W' if lon < 0 else 'E'}"

            base_name = source_file.stem
            new_filename = f"{lat_str}_{lon_str}_{base_name}.sub"
            dest_file = output_dir / new_filename

            # Copy .sub file
            shutil.copy(source_file, dest_file)

            # Create companion GPS JSON
            gps_json = {
                "type": "gps_coordinate",
                "source": "test_dataset",
                "city": city_name,
                "data": {
                    "latitude": lat,
                    "longitude": lon,
                    "accuracy": 5.0 + (i * 2),  # Vary accuracy
                    "altitude": 10.0 + (i * 5),
                    "timestamp": datetime.now(timezone.utc).isoformat()
                }
            }

            json_file = output_dir / new_filename.replace('.sub', '.json')
            with open(json_file, 'w') as f:
                json.dump(gps_json, f, indent=2)

            print(f"  ✅ {new_filename} → {lat:.4f}, {lon:.4f}")

            created_count += 1
            location_idx += 1

    # Create Wigle CSV manifest
    create_wigle_csv(output_dir)

    print(f"\n{'=' * 80}")
    print(f"✅ Created {created_count} GPS-tagged captures")
    print(f"📂 Output directory: {output_dir}")
    print(f"{'=' * 80}")

    return output_dir


def create_wigle_csv(output_dir):
    """Create Wigle CSV manifest for all captures"""

    csv_file = output_dir / "test_wigle_manifest.csv"

    # Get all .sub files
    sub_files = sorted(output_dir.glob("*.sub"))

    with open(csv_file, 'w') as f:
        # Header
        f.write("WigleWifi-1.4,appRelease=GigLez-Test-1.0,model=TestDataset\n")
        f.write("SignalHash,Frequency(MHz),Protocol,RSSI(dBm),Latitude,Longitude,Altitude,Accuracy,FirstSeen,LastSeen,CaptureCount,Interpolated,Confidence,BruceFile,SessionID\n")

        # Parse each file
        for sub_file in sub_files:
            # Parse .sub file for frequency/protocol
            frequency = 433.92  # Default
            protocol = "RAW"

            with open(sub_file, 'r') as sf:
                for line in sf:
                    if line.startswith('Frequency:'):
                        freq_hz = int(line.split(':')[1].strip())
                        frequency = freq_hz / 1e6
                    elif line.startswith('Protocol:'):
                        protocol = line.split(':')[1].strip()

            # Extract GPS from filename
            filename = sub_file.name
            parts = filename.split('_')

            if len(parts) >= 2:
                try:
                    # Parse latitude
                    lat_str = parts[0]
                    lat_val = float(lat_str[:-1])
                    if lat_str[-1] == 'S':
                        lat_val = -lat_val

                    # Parse longitude
                    lon_str = parts[1]
                    lon_val = float(lon_str[:-1])
                    if lon_str[-1] == 'W':
                        lon_val = -lon_val

                    # Write CSV row
                    hash_val = f"test_{sub_file.stem}"
                    timestamp = datetime.now(timezone.utc).isoformat()

                    f.write(f"{hash_val},{frequency:.3f},{protocol},-70,{lat_val},{lon_val},10.0,5.0,{timestamp},{timestamp},1,false,1.0,{filename},test_session\n")

                except Exception as e:
                    print(f"⚠️  Failed to parse GPS from {filename}: {e}")

    print(f"\n📄 Created Wigle CSV: {csv_file.name} ({len(sub_files)} captures)")


def main():
    """Create test dataset"""

    try:
        output_dir = create_gps_tagged_captures()

        print("\n📤 To upload all test data:")
        print(f"   python3 scripts/import_wardriving_data.py {output_dir}/")
        print("\n   OR")
        print(f"   python3 scripts/import_wardriving_data.py {output_dir}/test_wigle_manifest.csv")

        return 0

    except Exception as e:
        print(f"\n❌ Error: {e}")
        import traceback
        traceback.print_exc()
        return 1


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