225 lines
7.1 KiB
Python
225 lines
7.1 KiB
Python
import logging
|
|
import multiprocessing as mp
|
|
import socket
|
|
import struct
|
|
import subprocess
|
|
import threading
|
|
import time
|
|
from datetime import datetime
|
|
from typing import Callable, NamedTuple
|
|
|
|
import numpy as np
|
|
|
|
from ..config import config
|
|
from .csi_frame import CSI
|
|
from .protocols import CSICallback
|
|
|
|
Host = tuple[str, int]
|
|
|
|
|
|
class FeitHost:
|
|
def __init__(self, host: Host, command: str) -> None:
|
|
self.command = command
|
|
self.host = host
|
|
self.logger = logging.getLogger(
|
|
f"{__name__}.{self.__class__.__name__}-{self.host[0]}"
|
|
)
|
|
self.active = True
|
|
self.checker = threading.Thread(target=self.check_continuous)
|
|
self.checker.start()
|
|
|
|
def check_connection(self) -> bool:
|
|
feitcsi_status = subprocess.run(
|
|
f"ssh root@{self.host[0]} pgrep feitcsi",
|
|
check=False,
|
|
stdout=subprocess.DEVNULL,
|
|
shell=True,
|
|
)
|
|
return feitcsi_status.returncode == 0
|
|
|
|
def connect(self) -> None:
|
|
self.server = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
|
|
self.server.connect(self.host)
|
|
self.server.send(b"stop\n")
|
|
self.server.send(self.command.encode())
|
|
self.logger.info(f"Connected to {self.host}")
|
|
|
|
def check_continuous(self) -> None:
|
|
"""
|
|
Repeatedly check if the FeitCSI service is running
|
|
|
|
Because FeitCSI is using TCP, we will get no information if the service stops,
|
|
or the computer is not reachable. In order to make debugging easier, this checks
|
|
and logs continuously if the service is running.
|
|
"""
|
|
last_status = False
|
|
while self.active:
|
|
if not self.check_connection():
|
|
self.logger.error(f"FeitCSI is not running on {self.host[0]}")
|
|
last_status = False
|
|
else:
|
|
if not last_status:
|
|
self.connect()
|
|
last_status = True
|
|
time.sleep(1)
|
|
self.logger.info("Stopping host checker")
|
|
|
|
|
|
class FeitTransmitter(FeitHost):
|
|
def __init__(self) -> None:
|
|
command = (
|
|
f"feitcsi --frequency {config.central_freq} "
|
|
f"--channel-width {config.channel_width} "
|
|
f"--format {config.frame_format} "
|
|
f"--mode inject -s 1 --verbose "
|
|
f"--inject-delay {1_000_000 // config.sample_rate}"
|
|
)
|
|
super().__init__(config.transmit_host, command)
|
|
|
|
|
|
class FeitReceiver(FeitHost):
|
|
def __init__(self, host: Host) -> None:
|
|
command = (
|
|
f"feitcsi --frequency {config.central_freq} "
|
|
f"--channel-width {config.channel_width} "
|
|
f"--format {config.frame_format} "
|
|
f"--mode measure"
|
|
)
|
|
super().__init__(host, command)
|
|
|
|
def listen(self, queue: "mp.Queue[CSI]") -> None:
|
|
prev_time = datetime.now()
|
|
self.logger.info("Listening for CSI data")
|
|
while self.active:
|
|
while not hasattr(self, "server"):
|
|
time.sleep(0.1)
|
|
# This is the max size of a UDP packet. The size of the actual CSI
|
|
# packet will depend on the frame format and channel width, which
|
|
# changes the number of subcarriers
|
|
data = self.server.recv(65535)
|
|
try:
|
|
csidata = CSI(data)
|
|
self.logger.debug(
|
|
f"Received CSI data after {datetime.now() - prev_time}"
|
|
)
|
|
prev_time = datetime.now()
|
|
queue.put(csidata)
|
|
except struct.error:
|
|
self.logger.error("Failed to parse CSI data")
|
|
self.logger.info("Stopping CSI receiver")
|
|
|
|
|
|
class CSIProcessor:
|
|
def __init__(
|
|
self,
|
|
receiver_connections: dict[Host, "mp.Queue[CSI]"],
|
|
) -> None:
|
|
self.pending_data: dict[Host, tuple[datetime, CSI]] = {}
|
|
self.pending_data_lock = mp.Lock()
|
|
self.logger = logging.getLogger(f"{__name__}.{self.__class__.__name__}")
|
|
self.last_processed = datetime.now()
|
|
self.connections = receiver_connections
|
|
self.active = True
|
|
|
|
def add_data(self, host: Host, data: CSI) -> None:
|
|
if (
|
|
host in self.pending_data
|
|
and self.last_processed < self.pending_data[host][0]
|
|
):
|
|
self.logger.warning(
|
|
f"Skipping data from {host} at {self.pending_data[host][0]}"
|
|
)
|
|
|
|
with self.pending_data_lock:
|
|
self.pending_data[host] = (datetime.now(), data)
|
|
|
|
def is_ready(self) -> bool:
|
|
for host in config.receive_hosts:
|
|
if (
|
|
host not in self.pending_data
|
|
or self.pending_data[host][0] <= self.last_processed
|
|
):
|
|
return False
|
|
return True
|
|
|
|
def process_data(
|
|
self,
|
|
callback: CSICallback | None = None,
|
|
pre_merge_callback: Callable[[dict[Host, CSI]], None] | None = None,
|
|
) -> None:
|
|
self.last_processed = datetime.now()
|
|
if pre_merge_callback is not None:
|
|
pre_merge_callback(
|
|
{host: data[1] for host, data in self.pending_data.items()}
|
|
)
|
|
antenna_data = [
|
|
np.expand_dims(self.pending_data[ip][1].matrix[:, antenna], axis=2)
|
|
for ip, antenna in config.antennas.order
|
|
]
|
|
|
|
# We have data from all servers
|
|
all_data = np.concat(antenna_data, axis=1)
|
|
|
|
if callback:
|
|
callback(all_data)
|
|
|
|
def process_forever(
|
|
self,
|
|
callback: CSICallback | None = None,
|
|
pre_merge_callback: Callable[[dict[Host, CSI]], None] | None = None,
|
|
) -> None:
|
|
try:
|
|
while self.active:
|
|
for ip, queue in self.connections.items():
|
|
while not queue.empty():
|
|
self.add_data(ip, queue.get())
|
|
if self.is_ready():
|
|
self.process_data(callback, pre_merge_callback=pre_merge_callback)
|
|
else:
|
|
self.logger.debug("Not all data is ready")
|
|
time.sleep(0.0005)
|
|
except KeyboardInterrupt:
|
|
self.logger.info("Exiting CSI processing")
|
|
|
|
|
|
class Receiver(NamedTuple):
|
|
ip: Host
|
|
receiver: FeitReceiver
|
|
queue: "mp.Queue[CSI]"
|
|
|
|
|
|
def start_processing(
|
|
csi_callback: CSICallback | None = None,
|
|
pre_merge_callback: Callable[[dict[Host, CSI]], None] | None = None,
|
|
) -> None:
|
|
receivers = [
|
|
Receiver(ip, FeitReceiver(ip), mp.Queue(config.sample_rate))
|
|
for ip in config.receive_hosts
|
|
]
|
|
|
|
transmitter = FeitTransmitter()
|
|
|
|
# Start injecting CSI frames
|
|
processor = CSIProcessor({r.ip: r.queue for r in receivers})
|
|
|
|
receiver_processes = [
|
|
threading.Thread(target=r.receiver.listen, args=(r.queue,)) for r in receivers
|
|
]
|
|
|
|
for proc in receiver_processes:
|
|
proc.start()
|
|
|
|
processing_thread = mp.Process(
|
|
target=processor.process_forever, args=(csi_callback, pre_merge_callback)
|
|
)
|
|
processing_thread.start()
|
|
try:
|
|
while True:
|
|
time.sleep(100)
|
|
except KeyboardInterrupt:
|
|
for r in receivers:
|
|
r.receiver.active = False
|
|
transmitter.active = False
|
|
processor.active = False
|
|
return
|