414 lines
15 KiB
Python
414 lines
15 KiB
Python
"""Capture command for SDR devices."""
|
|
|
|
import os
|
|
|
|
import click
|
|
|
|
from utils.io import to_blue, to_npy, to_sigmf, to_wav
|
|
from utils.io.recording import generate_filename
|
|
from utils.view.view_signal_simple import view_simple_sig
|
|
|
|
from .common import (
|
|
echo_progress,
|
|
echo_verbose,
|
|
format_frequency,
|
|
format_sample_rate,
|
|
get_sdr_device,
|
|
load_yaml_config,
|
|
parse_frequency,
|
|
parse_metadata_args,
|
|
)
|
|
from .config import load_user_config
|
|
from .discover import (
|
|
find_bladerf_devices,
|
|
find_hackrf_devices,
|
|
find_pluto_devices,
|
|
find_rtlsdr_devices,
|
|
find_thinkrf_devices,
|
|
find_uhd_devices,
|
|
load_sdr_drivers,
|
|
)
|
|
|
|
|
|
def list_all_devices():
|
|
# Load drivers and collect all devices
|
|
load_sdr_drivers(verbose=False)
|
|
|
|
all_devices = []
|
|
all_devices.extend(find_uhd_devices())
|
|
all_devices.extend(find_pluto_devices())
|
|
all_devices.extend(find_hackrf_devices())
|
|
all_devices.extend(find_bladerf_devices())
|
|
all_devices.extend(find_rtlsdr_devices())
|
|
all_devices.extend(find_thinkrf_devices())
|
|
|
|
return all_devices
|
|
|
|
|
|
def auto_select_device(quiet: bool = False) -> str:
|
|
"""Auto-select device if only one is connected.
|
|
|
|
Args:
|
|
quiet: Suppress warning messages
|
|
|
|
Returns:
|
|
Device type string
|
|
|
|
Raises:
|
|
click.ClickException: If no devices or multiple devices found
|
|
"""
|
|
all_devices = list_all_devices()
|
|
|
|
if len(all_devices) == 0:
|
|
raise click.ClickException("No SDR devices found.\n" "Run 'utils discover' to see available devices.")
|
|
|
|
elif len(all_devices) == 1:
|
|
device = all_devices[0]
|
|
device_type = device.get("type", "Unknown").lower().replace("-", "").replace(" ", "")
|
|
|
|
# Map device type names to internal names
|
|
type_map = {
|
|
"plutosdr": "pluto",
|
|
"hackrf": "hackrf",
|
|
"hackrfone": "hackrf",
|
|
"bladerf": "bladerf",
|
|
"usrp": "usrp",
|
|
"b200": "usrp",
|
|
"b210": "usrp",
|
|
"rtlsdr": "rtlsdr",
|
|
"thinkrf": "thinkrf",
|
|
}
|
|
|
|
device_type = type_map.get(device_type, device_type)
|
|
|
|
if not quiet:
|
|
click.echo(
|
|
click.style("Warning: ", fg="yellow")
|
|
+ f"No device specified. Auto-detected {device.get('type', 'Unknown')}",
|
|
err=True,
|
|
)
|
|
click.echo(f"Use --device {device_type} to suppress this warning.\n", err=True)
|
|
|
|
return device_type
|
|
|
|
else:
|
|
device_list = "\n".join(f" - {d.get('type', 'Unknown')}" for d in all_devices)
|
|
raise click.ClickException(
|
|
f"Multiple devices found. Specify with --device\n\n"
|
|
f"Available devices:\n{device_list}\n\n"
|
|
f"Run 'utils discover' for more details."
|
|
)
|
|
|
|
|
|
def get_metadata_dict(config, metadata):
|
|
# Parse metadata - start with user config defaults
|
|
metadata_dict = config.get("metadata", {})
|
|
|
|
# Load user config and apply defaults
|
|
user_config = load_user_config()
|
|
|
|
# Apply user config metadata (if user config exists)
|
|
if user_config:
|
|
# Add standard metadata fields from config
|
|
for key in ["author", "organization", "project", "location", "testbed"]:
|
|
if key in user_config and key not in metadata_dict:
|
|
metadata_dict[key] = user_config[key]
|
|
|
|
# Add SigMF fields from config
|
|
if "sigmf" in user_config:
|
|
sigmf = user_config["sigmf"]
|
|
for key in ["license", "hw", "dataset"]:
|
|
if key in sigmf and key not in metadata_dict:
|
|
metadata_dict[key] = sigmf[key]
|
|
|
|
# CLI metadata overrides everything
|
|
if metadata:
|
|
metadata_dict.update(parse_metadata_args(metadata))
|
|
|
|
return metadata_dict
|
|
|
|
|
|
def save_visualization(recording, output_file: str, quiet: bool = False):
|
|
"""Save visualization of recording.
|
|
|
|
Args:
|
|
recording: Recording object
|
|
output_file: Path to save visualization (PNG)
|
|
quiet: Suppress progress messages
|
|
"""
|
|
# Generate image filename matching recording filename
|
|
base_name = os.path.splitext(output_file)[0]
|
|
if output_file.endswith(".sigmf-data"):
|
|
base_name = output_file.replace(".sigmf-data", "")
|
|
output_file = base_name + ".png"
|
|
|
|
try:
|
|
echo_progress(f"Generating visualization: {output_file}", quiet)
|
|
view_simple_sig(recording, output_path=output_file, saveplot=True, fast_mode=False, labels_mode=True)
|
|
except ImportError as e:
|
|
click.echo(click.style("Warning: ", fg="yellow") + f"Could not save visualization: {e}", err=True)
|
|
except Exception as e:
|
|
click.echo(click.style("Warning: ", fg="yellow") + f"Failed to save visualization: {e}", err=True)
|
|
|
|
|
|
def select_params(device, sample_rate, gain, bandwidth, quiet, verbose):
|
|
# Auto-select device if not specified
|
|
if device is None:
|
|
device = auto_select_device(quiet)
|
|
|
|
# Apply device-specific defaults (matching signal-testbed)
|
|
if sample_rate is None:
|
|
# Sample rate defaults based on signal-testbed hardware limits
|
|
device_sample_rates = {
|
|
"rtlsdr": 2.4e6, # RTL-SDR max is 3.2 MHz, use 2.4 MHz safe default
|
|
"thinkrf": 31.25e6, # ThinkRF decimation 4 (from 125 MS/s)
|
|
"pluto": 20e6, # PlutoSDR up to 61 MHz, 20 MHz safe
|
|
"hackrf": 20e6, # HackRF up to 20 MHz
|
|
"bladerf": 40e6, # BladeRF up to 61 MHz, 40 MHz safe
|
|
"usrp": 50e6, # USRP up to 200 MHz, 50 MHz default from signal-testbed
|
|
}
|
|
sample_rate = device_sample_rates.get(device, 20e6)
|
|
|
|
if gain is None:
|
|
# RX gain defaults (matching signal-testbed's 32 dB baseline, adjusted per device)
|
|
default_gains = {
|
|
"pluto": 32,
|
|
"hackrf": 32,
|
|
"bladerf": 32,
|
|
"usrp": 32,
|
|
"rtlsdr": 32, # RTL-SDR will auto-select closest valid gain
|
|
"thinkrf": 0, # ThinkRF uses attenuation, 0 = no attenuation
|
|
}
|
|
gain = default_gains.get(device, 32)
|
|
echo_verbose(f"Using default RX gain: {gain} dB for {device}", verbose)
|
|
|
|
if bandwidth is None:
|
|
# Bandwidth defaults (match sample rate for most devices)
|
|
device_bandwidths = {
|
|
"rtlsdr": None, # RTL-SDR doesn't support bandwidth setting
|
|
"thinkrf": None, # ThinkRF manages bandwidth internally
|
|
"pluto": sample_rate,
|
|
"hackrf": sample_rate,
|
|
"bladerf": sample_rate,
|
|
"usrp": sample_rate,
|
|
}
|
|
bandwidth = device_bandwidths.get(device)
|
|
|
|
return device, sample_rate, gain, bandwidth
|
|
|
|
|
|
def determine_output_format(output, output_format, output_dir):
|
|
# Determine output format and save
|
|
# If output specified, parse directory and filename
|
|
if output:
|
|
# Auto-detect format from extension if not specified
|
|
if output_format is None:
|
|
ext = os.path.splitext(output)[1].lower().lstrip(".")
|
|
if ext in ["sigmf", "sigmf-data"]:
|
|
output_format = "sigmf"
|
|
elif ext == "npy":
|
|
output_format = "npy"
|
|
elif ext == "wav":
|
|
output_format = "wav"
|
|
elif ext == "blue":
|
|
output_format = "blue"
|
|
else:
|
|
# Default to SigMF
|
|
output_format = "sigmf"
|
|
|
|
# Get output directory and filename from provided path
|
|
output_path_dir = os.path.dirname(output)
|
|
if output_path_dir:
|
|
output_dir = output_path_dir
|
|
output_filename = os.path.basename(output)
|
|
|
|
# Remove extension for formats that add it
|
|
if output_format == "sigmf":
|
|
output_filename = output_filename.replace(".sigmf-data", "").replace(".sigmf", "")
|
|
else:
|
|
# Use auto-generated filename based on timestamp and rec_id
|
|
output_filename = None # Will be auto-generated by save functions
|
|
if output_format is None:
|
|
output_format = "sigmf" # Default format
|
|
|
|
return output_format, output_filename, output_dir
|
|
|
|
|
|
# ============================================================================
|
|
# Main command
|
|
# ============================================================================
|
|
|
|
@click.command()
|
|
@click.argument("inputs", nargs=-1, required=True, type=click.Path(exists=True))
|
|
@click.argument("output", nargs=1, required=True, type=click.Path())
|
|
@click.option(
|
|
"--device",
|
|
"-d",
|
|
type=click.Choice(["pluto", "hackrf", "bladerf", "usrp", "rtlsdr", "thinkrf"]),
|
|
help="Device type",
|
|
)
|
|
@click.option("--ident", "-i", help="Device identifier (IP address or name=value, e.g., 192.168.2.1 or name=mypluto)")
|
|
@click.option(
|
|
"--config", "-c", "config_file", type=click.Path(exists=True), help="Load parameters from YAML config file"
|
|
)
|
|
@click.option(
|
|
"--sample-rate", "-s", type=float, default=None, help="Sample rate in Hz (e.g., 2e6) [default: device-specific]"
|
|
)
|
|
@click.option(
|
|
"--center-frequency",
|
|
"-f",
|
|
type=str,
|
|
default="2440M",
|
|
show_default=True,
|
|
help="Center frequency (e.g., 915e6, 2.4G)",
|
|
)
|
|
@click.option("--gain", "-g", type=float, help="RX gain in dB [default: device-specific]")
|
|
@click.option("--bandwidth", "-b", type=float, help="Bandwidth in Hz (if supported) [default: device-specific]")
|
|
@click.option("--num-samples", "-n", type=int, show_default=True, help="Number of samples to capture")
|
|
@click.option("--duration", "-t", type=float, help="Duration in seconds (alternative to --num-samples)")
|
|
@click.option("--output", "-o", help="Output filename (defaults to auto-generated with timestamp)")
|
|
@click.option("--output-dir", default="recordings", help="Output directory (default: recordings/)")
|
|
@click.option(
|
|
"--format",
|
|
"output_format",
|
|
type=click.Choice(["npy", "sigmf", "wav", "blue"]),
|
|
help="Output format (default: sigmf)",
|
|
)
|
|
@click.option("--save-image", is_flag=True, help="Save visualization PNG alongside recording")
|
|
@click.option("--metadata", "-m", multiple=True, help="Add custom metadata (KEY=VALUE)")
|
|
@click.option("--verbose", "-v", is_flag=True, help="Verbose output")
|
|
@click.option("--quiet", "-q", is_flag=True, help="Suppress progress output")
|
|
def capture(
|
|
device,
|
|
ident,
|
|
config_file,
|
|
sample_rate,
|
|
center_frequency,
|
|
gain,
|
|
bandwidth,
|
|
num_samples,
|
|
duration,
|
|
output,
|
|
output_dir,
|
|
output_format,
|
|
save_image,
|
|
metadata,
|
|
verbose,
|
|
quiet,
|
|
):
|
|
"""Capture IQ samples from SDR device and save to file.
|
|
|
|
\b
|
|
Examples:
|
|
utils capture -d hackrf -s 2e6 -f 2.44e6 -b 2e6
|
|
utils capture -d pluto -s 1e6 -f 2e9 -b 2e6 -n 50
|
|
|
|
"""
|
|
|
|
# Load config file if specified
|
|
config = {}
|
|
if config_file:
|
|
config = load_yaml_config(config_file)
|
|
echo_verbose(f"Loaded config from: {config_file}", verbose)
|
|
|
|
# Command-line args override config file
|
|
device = device or config.get("device")
|
|
ident = ident or config.get("ident") or config.get("serial") # Support legacy 'serial' in config
|
|
sample_rate = sample_rate or config.get("sample_rate")
|
|
center_frequency = center_frequency or config.get("center_frequency")
|
|
gain = gain or config.get("gain")
|
|
bandwidth = bandwidth or config.get("bandwidth")
|
|
num_samples = num_samples or config.get("num_samples")
|
|
duration = duration or config.get("duration")
|
|
output = output or config.get("output")
|
|
output_format = output_format or config.get("format")
|
|
|
|
# Parse metadata
|
|
metadata_dict = get_metadata_dict(config=config, metadata=metadata)
|
|
|
|
# Select parameters
|
|
device, sample_rate, gain, bandwidth = select_params(
|
|
device=device, sample_rate=sample_rate, gain=gain, bandwidth=bandwidth, quiet=quiet, verbose=verbose
|
|
)
|
|
|
|
# Parse frequency
|
|
center_freq_hz = parse_frequency(center_frequency)
|
|
|
|
# Calculate num_samples from duration if needed
|
|
if duration is not None and num_samples is None:
|
|
num_samples = int(duration * sample_rate)
|
|
echo_verbose(f"Duration {duration}s = {num_samples} samples at {format_sample_rate(sample_rate)}", verbose)
|
|
|
|
# Show capture parameters
|
|
echo_progress(f"Capturing from {device.upper()}...", quiet)
|
|
echo_progress(f"Sample rate: {format_sample_rate(sample_rate)}", quiet)
|
|
echo_progress(f"Center frequency: {format_frequency(center_freq_hz)}", quiet)
|
|
if gain is not None:
|
|
echo_progress(f"Gain: {gain} dB", quiet)
|
|
if bandwidth is not None:
|
|
echo_progress(f"Bandwidth: {format_sample_rate(bandwidth)}", quiet)
|
|
|
|
# Initialize device
|
|
echo_verbose("Initializing device...", verbose)
|
|
sdr = get_sdr_device(device, ident)
|
|
|
|
try:
|
|
# Initialize RX with parameters
|
|
echo_verbose("Initializing RX...", verbose)
|
|
sdr.init_rx(
|
|
sample_rate=sample_rate, center_frequency=center_freq_hz, gain=gain, channel=0 # Default to channel 0
|
|
)
|
|
|
|
# Set bandwidth if supported (after init_rx)
|
|
if bandwidth is not None and hasattr(sdr, "set_rx_bandwidth"):
|
|
sdr.set_rx_bandwidth(bandwidth)
|
|
|
|
# Capture
|
|
echo_progress(f"Capturing {num_samples} samples...", quiet)
|
|
recording = sdr.record(num_samples=num_samples)
|
|
|
|
echo_progress(
|
|
f"Captured {recording.data.shape[1] if len(recording.data.shape) > 1 else len(recording.data)} samples",
|
|
quiet,
|
|
)
|
|
|
|
# Add custom metadata to recording
|
|
if metadata_dict:
|
|
for key, value in metadata_dict.items():
|
|
recording.update_metadata(key, value)
|
|
|
|
output_format, output_filename, output_dir = determine_output_format(
|
|
output=output, output_format=output_format, output_dir=output_dir
|
|
)
|
|
echo_progress(f"Saving to {output_format.upper()} format...", quiet)
|
|
|
|
# Save recording (filenames with timestamp auto-generated if output_filename is None)
|
|
# All to_* functions handle directory creation internally
|
|
# Note: to_sigmf returns None, others return path
|
|
if output_format == "sigmf":
|
|
to_sigmf(recording, filename=output_filename, path=output_dir)
|
|
# Build path manually since to_sigmf doesn't return it
|
|
base_name = (
|
|
os.path.splitext(output_filename)[0] if output_filename else generate_filename(recording=recording)
|
|
)
|
|
saved_path = os.path.join(output_dir, f"{base_name}.sigmf-data")
|
|
elif output_format == "npy":
|
|
saved_path = to_npy(recording, filename=output_filename, path=output_dir)
|
|
elif output_format == "wav":
|
|
saved_path = to_wav(recording, filename=output_filename, path=output_dir)
|
|
elif output_format == "blue":
|
|
saved_path = to_blue(recording, filename=output_filename, path=output_dir)
|
|
|
|
echo_progress(f"Saved to: {saved_path}", quiet)
|
|
|
|
# Save visualization if requested
|
|
if save_image:
|
|
save_visualization(recording, saved_path, quiet)
|
|
|
|
finally:
|
|
# Clean up device
|
|
echo_verbose("Closing device...", verbose)
|
|
sdr.close()
|
|
|
|
echo_progress("Capture complete!", quiet)
|