Spaces:
Runtime error
Runtime error
| import math | |
| import os | |
| import re | |
| import platform | |
| import socket | |
| import subprocess | |
| import threading | |
| import time | |
| from collections import defaultdict | |
| from typing import Dict, List, Optional, Tuple | |
| from scapy.all import ARP, AsyncSniffer, Ether, IP, TCP, UDP, Raw, Scapy_Exception, conf, get_if_list, srp | |
| try: | |
| import psutil | |
| except ImportError: # pragma: no cover - local fallback keeps working if psutil is missing. | |
| psutil = None | |
| _WIFI_NAME_HINTS = ("wi-fi", "wifi", "wlan", "wireless", "802.11") | |
| def _safe_entropy(payload: bytes) -> float: | |
| if not payload: | |
| return 0.0 | |
| length = len(payload) | |
| counts = defaultdict(int) | |
| for byte in payload: | |
| counts[byte] += 1 | |
| entropy = 0.0 | |
| for count in counts.values(): | |
| probability = count / length | |
| entropy -= probability * math.log2(probability) | |
| return entropy | |
| def _is_loopback_interface(name: str) -> bool: | |
| lowered = name.lower() | |
| return lowered.startswith("lo") or lowered in {"loopback", "loopback pseudo-interface 1"} | |
| def _looks_like_wifi_interface(name: str) -> bool: | |
| lowered = name.lower() | |
| return any(hint in lowered for hint in _WIFI_NAME_HINTS) | |
| def _fallback_arp_cache() -> List[Dict[str, str]]: | |
| try: | |
| output = subprocess.check_output(["arp", "-a"], text=True, stderr=subprocess.DEVNULL) | |
| except Exception: | |
| return [] | |
| devices = [] | |
| pattern = re.compile(r"(?P<ip>\d+\.\d+\.\d+\.\d+)\s+(?P<mac>(?:[0-9a-f]{2}-){5}[0-9a-f]{2})", re.IGNORECASE) | |
| for match in pattern.finditer(output): | |
| devices.append( | |
| { | |
| "ip": match.group("ip"), | |
| "mac": match.group("mac").replace("-", ":").lower(), | |
| } | |
| ) | |
| return devices | |
| def _select_wifi_interface(preferred: Optional[str] = None) -> Optional[str]: | |
| if preferred: | |
| return preferred | |
| if psutil is not None: | |
| stats = psutil.net_if_stats() | |
| for name, stat in stats.items(): | |
| if stat.isup and not _is_loopback_interface(name) and _looks_like_wifi_interface(name): | |
| return name | |
| addrs = psutil.net_if_addrs() | |
| for name, stat in stats.items(): | |
| if stat.isup and not _is_loopback_interface(name) and any(addr.family == socket.AF_INET for addr in addrs.get(name, [])): | |
| return name | |
| for name in get_if_list(): | |
| if not _is_loopback_interface(name) and _looks_like_wifi_interface(name): | |
| return name | |
| for name in get_if_list(): | |
| if not _is_loopback_interface(name): | |
| return name | |
| return None | |
| def _interface_ipv4_address(interface: Optional[str]) -> Optional[str]: | |
| if not interface or psutil is None: | |
| return None | |
| try: | |
| for addr in psutil.net_if_addrs().get(interface, []): | |
| if addr.family == socket.AF_INET: | |
| return addr.address | |
| except Exception: | |
| return None | |
| return None | |
| class WiFiTelemetrySource: | |
| """Builds a live network snapshot from the WiFi adapter and ARP cache.""" | |
| def __init__(self): | |
| self._lock = threading.Lock() | |
| self._iface: Optional[str] = None | |
| self._last_counters = None | |
| self._last_sample_time: Optional[float] = None | |
| self._latest_features: List[Dict[str, object]] = [] | |
| self.last_error: Optional[str] = None | |
| self.is_running = False | |
| def start(self, iface: Optional[str] = None): | |
| with self._lock: | |
| self._iface = _select_wifi_interface(iface) | |
| self._last_counters = None | |
| self._last_sample_time = None | |
| self._latest_features = [] | |
| self.last_error = None | |
| self.is_running = True | |
| return self._iface | |
| def stop(self): | |
| with self._lock: | |
| self._last_counters = None | |
| self._last_sample_time = None | |
| self._latest_features = [] | |
| self.is_running = False | |
| def _build_features( | |
| self, | |
| arp_devices: List[Dict[str, str]], | |
| elapsed: float, | |
| packet_total: float, | |
| byte_total: float, | |
| interface: Optional[str], | |
| local_ip: Optional[str], | |
| ) -> List[Dict[str, object]]: | |
| devices = arp_devices or [{"ip": local_ip or "0.0.0.0", "mac": "local"}] | |
| per_device_count = max(len(devices), 1) | |
| packet_rate = packet_total / max(elapsed, 1.0) | |
| byte_rate = byte_total / max(elapsed, 1.0) | |
| packet_size = byte_total / max(packet_total, 1.0) if byte_total > 0 else 0.0 | |
| features = [] | |
| for device in devices: | |
| features.append( | |
| { | |
| "ip": device["ip"], | |
| "dst_ip": local_ip or device["ip"], | |
| "protocol": "WIFI", | |
| "protocol_num": 2, | |
| "port_dst": 0, | |
| "packets": int(max(round(packet_total / per_device_count), 0)), | |
| "packet_size": packet_size, | |
| "packet_rate": packet_rate / per_device_count if per_device_count else packet_rate, | |
| "byte_rate": byte_rate / per_device_count if per_device_count else byte_rate, | |
| "avg_pkt_size": packet_size, | |
| "tcp_flags": 0, | |
| "flow_duration": max(elapsed, 0.001), | |
| "entropy": 0.0, | |
| "ttl": 64.0, | |
| "raw_stats": { | |
| "source": "wifi_telemetry", | |
| "interface": interface, | |
| "local_ip": local_ip, | |
| "mac": device.get("mac"), | |
| "arp_devices": len(devices), | |
| }, | |
| } | |
| ) | |
| return features | |
| def get_latest_features(self): | |
| if psutil is None: | |
| self.last_error = "psutil is unavailable" | |
| return [] | |
| now = time.time() | |
| interface = _select_wifi_interface(self._iface) | |
| if interface is None: | |
| self.last_error = "No active WiFi-like interface found" | |
| return [] | |
| counters = psutil.net_io_counters(pernic=True).get(interface) | |
| if counters is None: | |
| self.last_error = f"No counters available for interface {interface}" | |
| return [] | |
| arp_devices = _fallback_arp_cache() | |
| local_ip = _interface_ipv4_address(interface) | |
| with self._lock: | |
| if self._last_counters is None: | |
| self._last_counters = counters | |
| self._last_sample_time = now | |
| self._latest_features = self._build_features( | |
| arp_devices=arp_devices, | |
| elapsed=1.0, | |
| packet_total=0.0, | |
| byte_total=0.0, | |
| interface=interface, | |
| local_ip=local_ip, | |
| ) | |
| self.last_error = None | |
| return list(self._latest_features) | |
| elapsed = max(now - (self._last_sample_time or now), 1.0) | |
| packet_total = max( | |
| (counters.packets_recv + counters.packets_sent) | |
| - (self._last_counters.packets_recv + self._last_counters.packets_sent), | |
| 0, | |
| ) | |
| byte_total = max( | |
| (counters.bytes_recv + counters.bytes_sent) | |
| - (self._last_counters.bytes_recv + self._last_counters.bytes_sent), | |
| 0, | |
| ) | |
| if packet_total == 0 and byte_total > 0: | |
| packet_total = max(int(byte_total / 900), 1) | |
| self._last_counters = counters | |
| self._last_sample_time = now | |
| self._latest_features = self._build_features( | |
| arp_devices=arp_devices, | |
| elapsed=elapsed, | |
| packet_total=float(packet_total), | |
| byte_total=float(byte_total), | |
| interface=interface, | |
| local_ip=local_ip, | |
| ) | |
| self.last_error = None | |
| return list(self._latest_features) | |
| def arp_scan(self, subnet: Optional[str] = None) -> List[Dict[str, str]]: | |
| devices = _fallback_arp_cache() | |
| if devices: | |
| return devices | |
| return [] | |
| def status(self) -> Dict[str, object]: | |
| with self._lock: | |
| return { | |
| "mode": "wifi_telemetry", | |
| "interface": self._iface or _select_wifi_interface(), | |
| "running": self.is_running, | |
| "last_error": self.last_error, | |
| } | |
| class LiveCapture: | |
| def __init__(self): | |
| self._lock = threading.Lock() | |
| self._sniffer: Optional[AsyncSniffer] = None | |
| self._sniff_iface: Optional[str] = None | |
| self.last_error: Optional[str] = None | |
| self.flow_stats = defaultdict( | |
| lambda: { | |
| "packets": 0, | |
| "bytes": 0, | |
| "first_seen": time.time(), | |
| "last_seen": time.time(), | |
| "protocol": "UNKNOWN", | |
| "dst_port": 0, | |
| "tcp_flags": 0, | |
| "ttl_sum": 0, | |
| "ttl_count": 0, | |
| "payload_entropy_sum": 0.0, | |
| "payload_entropy_count": 0, | |
| } | |
| ) | |
| def process_packet(self, pkt): | |
| if IP not in pkt: | |
| return | |
| src_ip = pkt[IP].src | |
| dst_ip = pkt[IP].dst | |
| protocol = "OTHER" | |
| protocol_num = 2 | |
| dst_port = 0 | |
| tcp_flags = 0 | |
| if TCP in pkt: | |
| protocol = "TCP" | |
| protocol_num = 0 | |
| dst_port = int(pkt[TCP].dport) | |
| tcp_flags = int(pkt[TCP].flags) | |
| elif UDP in pkt: | |
| protocol = "UDP" | |
| protocol_num = 1 | |
| dst_port = int(pkt[UDP].dport) | |
| key = (src_ip, dst_ip, protocol, dst_port) | |
| payload = bytes(pkt[Raw].load) if Raw in pkt else b"" | |
| entropy = _safe_entropy(payload) | |
| with self._lock: | |
| flow = self.flow_stats[key] | |
| now = time.time() | |
| if flow["packets"] == 0: | |
| flow["first_seen"] = now | |
| flow["last_seen"] = now | |
| flow["packets"] += 1 | |
| flow["bytes"] += len(pkt) | |
| flow["protocol"] = protocol | |
| flow["protocol_num"] = protocol_num | |
| flow["dst_port"] = dst_port | |
| flow["tcp_flags"] |= tcp_flags | |
| flow["ttl_sum"] += int(pkt[IP].ttl) | |
| flow["ttl_count"] += 1 | |
| flow["payload_entropy_sum"] += entropy | |
| flow["payload_entropy_count"] += 1 | |
| def start_capture(self, iface: Optional[str] = None): | |
| if self._sniffer and self._sniffer.running: | |
| return | |
| self._sniff_iface = _select_wifi_interface(iface) | |
| self.last_error = None | |
| if platform.system() == "Windows" and not getattr(conf, "use_pcap", False): | |
| self.last_error = "Npcap/WinPcap is not available on this machine" | |
| raise RuntimeError(self.last_error) | |
| try: | |
| self._sniffer = AsyncSniffer(iface=self._sniff_iface, prn=self.process_packet, store=False) | |
| self._sniffer.start() | |
| except (PermissionError, OSError, RuntimeError, Scapy_Exception) as exc: | |
| self.last_error = str(exc) | |
| self._sniffer = None | |
| raise | |
| def stop_capture(self): | |
| if self._sniffer and self._sniffer.running: | |
| try: | |
| self._sniffer.stop() | |
| except RuntimeError: | |
| pass | |
| finally: | |
| self._sniffer = None | |
| def arp_scan(self, subnet: Optional[str] = None) -> List[Dict[str, str]]: | |
| subnet = subnet or os.getenv("NETWOKIE_SUBNET", "192.168.1.0/24") | |
| try: | |
| ans, _ = srp( | |
| Ether(dst="ff:ff:ff:ff:ff:ff") / ARP(pdst=subnet), | |
| timeout=2, | |
| verbose=False, | |
| ) | |
| except (PermissionError, OSError, RuntimeError, Scapy_Exception) as exc: | |
| self.last_error = str(exc) | |
| return _fallback_arp_cache() | |
| return [{"ip": r[1].psrc, "mac": r[1].hwsrc} for r in ans] | |
| def get_latest_features(self): | |
| current_time = time.time() | |
| features = [] | |
| with self._lock: | |
| snapshot = list(self.flow_stats.items()) | |
| for (src_ip, dst_ip, protocol, dst_port), stats in snapshot: | |
| duration = max(current_time - stats["first_seen"], 0.001) | |
| packets = max(stats["packets"], 1) | |
| features.append( | |
| { | |
| "ip": src_ip, | |
| "dst_ip": dst_ip, | |
| "protocol": protocol, | |
| "protocol_num": stats["protocol_num"], | |
| "port_dst": dst_port, | |
| "packets": stats["packets"], | |
| "packet_size": stats["bytes"] / packets, | |
| "packet_rate": stats["packets"] / duration, | |
| "byte_rate": stats["bytes"] / duration, | |
| "avg_pkt_size": stats["bytes"] / packets, | |
| "tcp_flags": stats["tcp_flags"], | |
| "flow_duration": duration, | |
| "entropy": ( | |
| stats["payload_entropy_sum"] / stats["payload_entropy_count"] | |
| if stats["payload_entropy_count"] | |
| else 0.0 | |
| ), | |
| "ttl": ( | |
| stats["ttl_sum"] / stats["ttl_count"] | |
| if stats["ttl_count"] | |
| else 0.0 | |
| ), | |
| "raw_stats": stats, | |
| } | |
| ) | |
| return features | |
| def clear_stale_flows(self, timeout_seconds=60): | |
| current_time = time.time() | |
| with self._lock: | |
| stale_keys = [ | |
| key | |
| for key, stats in self.flow_stats.items() | |
| if current_time - stats["last_seen"] > timeout_seconds | |
| ] | |
| for key in stale_keys: | |
| del self.flow_stats[key] | |