mirror of
https://github.com/infinition/Bjorn.git
synced 2026-03-15 17:01:58 +00:00
Add Loki and Sentinel utility classes for web API endpoints
- Implemented LokiUtils class with GET and POST endpoints for managing scripts, jobs, and payloads. - Added SentinelUtils class with GET and POST endpoints for managing events, rules, devices, and notifications. - Both classes include error handling and JSON response formatting.
This commit is contained in:
585
bifrost/__init__.py
Normal file
585
bifrost/__init__.py
Normal file
@@ -0,0 +1,585 @@
|
||||
"""
|
||||
Bifrost — Pwnagotchi-compatible WiFi recon engine for Bjorn.
|
||||
Runs as a daemon thread alongside MANUAL/AUTO/AI modes.
|
||||
"""
|
||||
import os
|
||||
import time
|
||||
import subprocess
|
||||
import threading
|
||||
import logging
|
||||
|
||||
from logger import Logger
|
||||
|
||||
logger = Logger(name="bifrost", level=logging.DEBUG)
|
||||
|
||||
|
||||
class BifrostEngine:
|
||||
"""Main Bifrost lifecycle manager.
|
||||
|
||||
Manages the bettercap subprocess and BifrostAgent daemon loop.
|
||||
Pattern follows SentinelEngine (sentinel.py).
|
||||
"""
|
||||
|
||||
def __init__(self, shared_data):
|
||||
self.shared_data = shared_data
|
||||
self._thread = None
|
||||
self._stop_event = threading.Event()
|
||||
self._running = False
|
||||
self._bettercap_proc = None
|
||||
self._monitor_torn_down = False
|
||||
self._monitor_failed = False
|
||||
self.agent = None
|
||||
|
||||
@property
|
||||
def enabled(self):
|
||||
return bool(self.shared_data.config.get('bifrost_enabled', False))
|
||||
|
||||
def start(self):
|
||||
"""Start the Bifrost engine (bettercap + agent loop)."""
|
||||
if self._running:
|
||||
logger.warning("Bifrost already running")
|
||||
return
|
||||
|
||||
# Wait for any previous thread to finish before re-starting
|
||||
if self._thread and self._thread.is_alive():
|
||||
logger.warning("Previous Bifrost thread still running — waiting ...")
|
||||
self._stop_event.set()
|
||||
self._thread.join(timeout=15)
|
||||
|
||||
logger.info("Starting Bifrost engine ...")
|
||||
self._stop_event.clear()
|
||||
self._running = True
|
||||
self._monitor_failed = False
|
||||
self._monitor_torn_down = False
|
||||
|
||||
self._thread = threading.Thread(
|
||||
target=self._loop, daemon=True, name="BifrostEngine"
|
||||
)
|
||||
self._thread.start()
|
||||
|
||||
def stop(self):
|
||||
"""Stop the Bifrost engine gracefully.
|
||||
|
||||
Signals the daemon loop to exit, then waits for it to finish.
|
||||
The loop's finally block handles bettercap shutdown and monitor teardown.
|
||||
"""
|
||||
if not self._running:
|
||||
return
|
||||
|
||||
logger.info("Stopping Bifrost engine ...")
|
||||
self._stop_event.set()
|
||||
self._running = False
|
||||
|
||||
if self._thread and self._thread.is_alive():
|
||||
self._thread.join(timeout=15)
|
||||
self._thread = None
|
||||
self.agent = None
|
||||
|
||||
# Safety net: teardown is idempotent, so this is a no-op if
|
||||
# _loop()'s finally already ran it.
|
||||
self._stop_bettercap()
|
||||
self._teardown_monitor_mode()
|
||||
logger.info("Bifrost engine stopped")
|
||||
|
||||
def _loop(self):
|
||||
"""Main daemon loop — setup monitor mode, start bettercap, create agent, run recon cycle."""
|
||||
try:
|
||||
# Install compatibility shim for pwnagotchi plugins
|
||||
from bifrost import plugins as bfplugins
|
||||
from bifrost.compat import install_shim
|
||||
install_shim(self.shared_data, bfplugins)
|
||||
|
||||
# Setup monitor mode on the WiFi interface
|
||||
self._setup_monitor_mode()
|
||||
|
||||
if self._monitor_failed:
|
||||
logger.error(
|
||||
"Monitor mode setup failed — Bifrost cannot operate without monitor "
|
||||
"mode. For Broadcom chips (Pi Zero W/2W), install nexmon: "
|
||||
"https://github.com/seemoo-lab/nexmon — "
|
||||
"Or use an external USB WiFi adapter with monitor mode support.")
|
||||
# Teardown first (restores network services) BEFORE switching mode,
|
||||
# so the orchestrator doesn't start scanning on a dead network.
|
||||
self._teardown_monitor_mode()
|
||||
self._running = False
|
||||
# Now switch mode back to AUTO — the network should be restored.
|
||||
# We set the flag directly FIRST (bypass setter to avoid re-stopping),
|
||||
# then ensure manual_mode/ai_mode are cleared so getter returns AUTO.
|
||||
try:
|
||||
self.shared_data.config["bifrost_enabled"] = False
|
||||
self.shared_data.config["manual_mode"] = False
|
||||
self.shared_data.config["ai_mode"] = False
|
||||
self.shared_data.manual_mode = False
|
||||
self.shared_data.ai_mode = False
|
||||
self.shared_data.invalidate_config_cache()
|
||||
logger.info("Bifrost auto-disabled due to monitor mode failure — mode: AUTO")
|
||||
except Exception:
|
||||
pass
|
||||
return
|
||||
|
||||
# Start bettercap
|
||||
self._start_bettercap()
|
||||
self._stop_event.wait(3) # Give bettercap time to initialize
|
||||
if self._stop_event.is_set():
|
||||
return
|
||||
|
||||
# Create agent (pass stop_event so its threads exit cleanly)
|
||||
from bifrost.agent import BifrostAgent
|
||||
self.agent = BifrostAgent(self.shared_data, stop_event=self._stop_event)
|
||||
|
||||
# Load plugins
|
||||
bfplugins.load(self.shared_data.config)
|
||||
|
||||
# Initialize agent
|
||||
self.agent.start()
|
||||
|
||||
logger.info("Bifrost agent started — entering recon cycle")
|
||||
|
||||
# Main recon loop (port of do_auto_mode from pwnagotchi)
|
||||
while not self._stop_event.is_set():
|
||||
try:
|
||||
# Full spectrum scan
|
||||
self.agent.recon()
|
||||
|
||||
if self._stop_event.is_set():
|
||||
break
|
||||
|
||||
# Get APs grouped by channel
|
||||
channels = self.agent.get_access_points_by_channel()
|
||||
|
||||
# For each channel
|
||||
for ch, aps in channels:
|
||||
if self._stop_event.is_set():
|
||||
break
|
||||
|
||||
self.agent.set_channel(ch)
|
||||
|
||||
# For each AP on this channel
|
||||
for ap in aps:
|
||||
if self._stop_event.is_set():
|
||||
break
|
||||
|
||||
# Send association frame for PMKID
|
||||
self.agent.associate(ap)
|
||||
|
||||
# Deauth all clients for full handshake
|
||||
for sta in ap.get('clients', []):
|
||||
if self._stop_event.is_set():
|
||||
break
|
||||
self.agent.deauth(ap, sta)
|
||||
|
||||
if not self._stop_event.is_set():
|
||||
self.agent.next_epoch()
|
||||
|
||||
except Exception as e:
|
||||
if 'wifi.interface not set' in str(e):
|
||||
logger.error("WiFi interface lost: %s", e)
|
||||
self._stop_event.wait(60)
|
||||
if not self._stop_event.is_set():
|
||||
self.agent.next_epoch()
|
||||
else:
|
||||
logger.error("Recon loop error: %s", e)
|
||||
self._stop_event.wait(5)
|
||||
|
||||
except Exception as e:
|
||||
logger.error("Bifrost engine fatal error: %s", e)
|
||||
finally:
|
||||
from bifrost import plugins as bfplugins
|
||||
bfplugins.shutdown()
|
||||
self._stop_bettercap()
|
||||
self._teardown_monitor_mode()
|
||||
self._running = False
|
||||
|
||||
# ── Monitor mode management ─────────────────────────
|
||||
|
||||
# ── Nexmon helpers ────────────────────────────────────
|
||||
|
||||
@staticmethod
|
||||
def _has_nexmon():
|
||||
"""Check if nexmon firmware patches are installed."""
|
||||
import shutil
|
||||
if not shutil.which('nexutil'):
|
||||
return False
|
||||
# Verify patched firmware via dmesg
|
||||
try:
|
||||
r = subprocess.run(
|
||||
['dmesg'], capture_output=True, text=True, timeout=5)
|
||||
if 'nexmon' in r.stdout.lower():
|
||||
return True
|
||||
except Exception:
|
||||
pass
|
||||
# nexutil exists — assume usable even without dmesg confirmation
|
||||
return True
|
||||
|
||||
@staticmethod
|
||||
def _is_brcmfmac(iface):
|
||||
"""Check if the interface uses the brcmfmac driver (Broadcom)."""
|
||||
driver_path = '/sys/class/net/%s/device/driver' % iface
|
||||
try:
|
||||
real = os.path.realpath(driver_path)
|
||||
return 'brcmfmac' in real
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
def _detect_phy(self, iface):
|
||||
"""Detect the phy name for a given interface (e.g. 'phy0')."""
|
||||
try:
|
||||
r = subprocess.run(
|
||||
['iw', 'dev', iface, 'info'],
|
||||
capture_output=True, text=True, timeout=5)
|
||||
for line in r.stdout.splitlines():
|
||||
if 'wiphy' in line:
|
||||
idx = line.strip().split()[-1]
|
||||
return 'phy%s' % idx
|
||||
except Exception:
|
||||
pass
|
||||
return 'phy0'
|
||||
|
||||
def _setup_monitor_mode(self):
|
||||
"""Put the WiFi interface into monitor mode.
|
||||
|
||||
Strategy order:
|
||||
1. Nexmon — for Broadcom brcmfmac chips (Pi Zero W / Pi Zero 2 W)
|
||||
Uses: iw phy <phy> interface add mon0 type monitor + nexutil -m2
|
||||
2. airmon-ng — for chipsets with proper driver support (Atheros, Realtek, etc.)
|
||||
3. iw — direct fallback for other drivers
|
||||
"""
|
||||
self._monitor_torn_down = False
|
||||
self._nexmon_used = False
|
||||
cfg = self.shared_data.config
|
||||
iface = cfg.get('bifrost_iface', 'wlan0mon')
|
||||
|
||||
# If configured iface already ends with 'mon', derive the base name
|
||||
if iface.endswith('mon'):
|
||||
base_iface = iface[:-3] # e.g. 'wlan0mon' -> 'wlan0'
|
||||
else:
|
||||
base_iface = iface
|
||||
|
||||
# Store original interface name for teardown
|
||||
self._base_iface = base_iface
|
||||
self._mon_iface = iface
|
||||
|
||||
# Check if a monitor interface already exists
|
||||
if iface != base_iface and self._iface_exists(iface):
|
||||
logger.info("Monitor interface %s already exists", iface)
|
||||
return
|
||||
|
||||
# ── Strategy 1: Nexmon (Broadcom brcmfmac) ────────────────
|
||||
if self._is_brcmfmac(base_iface):
|
||||
logger.info("Broadcom brcmfmac chip detected on %s", base_iface)
|
||||
if self._has_nexmon():
|
||||
if self._setup_nexmon(base_iface, cfg):
|
||||
return
|
||||
# nexmon setup failed — don't try other strategies, they won't work either
|
||||
self._monitor_failed = True
|
||||
return
|
||||
else:
|
||||
logger.error(
|
||||
"Broadcom brcmfmac chip requires nexmon firmware patches for "
|
||||
"monitor mode. Install nexmon manually using install_nexmon.sh "
|
||||
"or visit: https://github.com/seemoo-lab/nexmon")
|
||||
self._monitor_failed = True
|
||||
return
|
||||
|
||||
# ── Strategy 2: airmon-ng (Atheros, Realtek, etc.) ────────
|
||||
airmon_ok = False
|
||||
try:
|
||||
logger.info("Killing interfering processes ...")
|
||||
subprocess.run(
|
||||
['airmon-ng', 'check', 'kill'],
|
||||
stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
|
||||
timeout=15,
|
||||
)
|
||||
logger.info("Starting monitor mode: airmon-ng start %s", base_iface)
|
||||
result = subprocess.run(
|
||||
['airmon-ng', 'start', base_iface],
|
||||
capture_output=True, text=True, timeout=30,
|
||||
)
|
||||
combined = (result.stdout + result.stderr).strip()
|
||||
logger.info("airmon-ng output: %s", combined)
|
||||
|
||||
if 'Operation not supported' in combined or 'command failed' in combined:
|
||||
logger.warning("airmon-ng failed: %s", combined)
|
||||
else:
|
||||
# airmon-ng may rename the interface (wlan0 -> wlan0mon)
|
||||
if self._iface_exists(iface):
|
||||
logger.info("Monitor mode active: %s", iface)
|
||||
airmon_ok = True
|
||||
elif self._iface_exists(base_iface):
|
||||
logger.info("Interface %s is now in monitor mode (no rename)", base_iface)
|
||||
cfg['bifrost_iface'] = base_iface
|
||||
self._mon_iface = base_iface
|
||||
airmon_ok = True
|
||||
|
||||
if airmon_ok:
|
||||
return
|
||||
except FileNotFoundError:
|
||||
logger.warning("airmon-ng not found, trying iw fallback ...")
|
||||
except Exception as e:
|
||||
logger.warning("airmon-ng failed: %s, trying iw fallback ...", e)
|
||||
|
||||
# ── Strategy 3: iw (direct fallback) ──────────────────────
|
||||
try:
|
||||
subprocess.run(
|
||||
['ip', 'link', 'set', base_iface, 'down'],
|
||||
stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, timeout=10,
|
||||
)
|
||||
result = subprocess.run(
|
||||
['iw', 'dev', base_iface, 'set', 'type', 'monitor'],
|
||||
capture_output=True, text=True, timeout=10,
|
||||
)
|
||||
if result.returncode != 0:
|
||||
err = result.stderr.strip()
|
||||
logger.error("iw set monitor failed (rc=%d): %s", result.returncode, err)
|
||||
self._monitor_failed = True
|
||||
subprocess.run(
|
||||
['ip', 'link', 'set', base_iface, 'up'],
|
||||
stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, timeout=10,
|
||||
)
|
||||
return
|
||||
subprocess.run(
|
||||
['ip', 'link', 'set', base_iface, 'up'],
|
||||
stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, timeout=10,
|
||||
)
|
||||
logger.info("Monitor mode set via iw on %s", base_iface)
|
||||
cfg['bifrost_iface'] = base_iface
|
||||
self._mon_iface = base_iface
|
||||
except Exception as e:
|
||||
logger.error("Failed to set monitor mode: %s", e)
|
||||
self._monitor_failed = True
|
||||
|
||||
def _setup_nexmon(self, base_iface, cfg):
|
||||
"""Enable monitor mode using nexmon (for Broadcom brcmfmac chips).
|
||||
|
||||
Creates a separate monitor interface (mon0) so wlan0 can potentially
|
||||
remain usable for management traffic (like pwnagotchi does).
|
||||
|
||||
Returns True on success, False on failure.
|
||||
"""
|
||||
mon_iface = 'mon0'
|
||||
phy = self._detect_phy(base_iface)
|
||||
logger.info("Nexmon: setting up monitor mode on %s (phy=%s)", base_iface, phy)
|
||||
|
||||
try:
|
||||
# Kill interfering services (same as pwnagotchi)
|
||||
for svc in ('wpa_supplicant', 'NetworkManager', 'dhcpcd'):
|
||||
subprocess.run(
|
||||
['systemctl', 'stop', svc],
|
||||
stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, timeout=10,
|
||||
)
|
||||
|
||||
# Remove old mon0 if it exists
|
||||
if self._iface_exists(mon_iface):
|
||||
subprocess.run(
|
||||
['iw', 'dev', mon_iface, 'del'],
|
||||
stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, timeout=5,
|
||||
)
|
||||
|
||||
# Create monitor interface via iw phy
|
||||
result = subprocess.run(
|
||||
['iw', 'phy', phy, 'interface', 'add', mon_iface, 'type', 'monitor'],
|
||||
capture_output=True, text=True, timeout=10,
|
||||
)
|
||||
if result.returncode != 0:
|
||||
logger.error("Failed to create %s: %s", mon_iface, result.stderr.strip())
|
||||
return False
|
||||
|
||||
# Bring monitor interface up
|
||||
subprocess.run(
|
||||
['ifconfig', mon_iface, 'up'],
|
||||
stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, timeout=10,
|
||||
)
|
||||
|
||||
# Enable monitor mode with radiotap headers via nexutil
|
||||
result = subprocess.run(
|
||||
['nexutil', '-m2'],
|
||||
capture_output=True, text=True, timeout=10,
|
||||
)
|
||||
if result.returncode != 0:
|
||||
logger.warning("nexutil -m2 returned rc=%d: %s", result.returncode, result.stderr.strip())
|
||||
|
||||
# Verify
|
||||
verify = subprocess.run(
|
||||
['nexutil', '-m'],
|
||||
capture_output=True, text=True, timeout=5,
|
||||
)
|
||||
mode_val = verify.stdout.strip()
|
||||
logger.info("nexutil -m reports: %s", mode_val)
|
||||
|
||||
if not self._iface_exists(mon_iface):
|
||||
logger.error("Monitor interface %s not created", mon_iface)
|
||||
return False
|
||||
|
||||
# Success — update config to use mon0
|
||||
cfg['bifrost_iface'] = mon_iface
|
||||
self._mon_iface = mon_iface
|
||||
self._nexmon_used = True
|
||||
logger.info("Nexmon monitor mode active on %s (phy=%s)", mon_iface, phy)
|
||||
return True
|
||||
|
||||
except FileNotFoundError as e:
|
||||
logger.error("Required tool not found: %s", e)
|
||||
return False
|
||||
except Exception as e:
|
||||
logger.error("Nexmon setup error: %s", e)
|
||||
return False
|
||||
|
||||
def _teardown_monitor_mode(self):
|
||||
"""Restore the WiFi interface to managed mode (idempotent)."""
|
||||
if self._monitor_torn_down:
|
||||
return
|
||||
base_iface = getattr(self, '_base_iface', None)
|
||||
mon_iface = getattr(self, '_mon_iface', None)
|
||||
if not base_iface:
|
||||
return
|
||||
self._monitor_torn_down = True
|
||||
|
||||
logger.info("Restoring managed mode for %s ...", base_iface)
|
||||
|
||||
if getattr(self, '_nexmon_used', False):
|
||||
# ── Nexmon teardown ──
|
||||
try:
|
||||
subprocess.run(
|
||||
['nexutil', '-m0'],
|
||||
stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, timeout=5,
|
||||
)
|
||||
logger.info("Nexmon monitor mode disabled (nexutil -m0)")
|
||||
except Exception:
|
||||
pass
|
||||
# Remove the mon0 interface
|
||||
if mon_iface and mon_iface != base_iface and self._iface_exists(mon_iface):
|
||||
try:
|
||||
subprocess.run(
|
||||
['iw', 'dev', mon_iface, 'del'],
|
||||
stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, timeout=5,
|
||||
)
|
||||
logger.info("Removed monitor interface %s", mon_iface)
|
||||
except Exception:
|
||||
pass
|
||||
else:
|
||||
# ── airmon-ng / iw teardown ──
|
||||
try:
|
||||
iface_to_stop = mon_iface or base_iface
|
||||
subprocess.run(
|
||||
['airmon-ng', 'stop', iface_to_stop],
|
||||
stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
|
||||
timeout=15,
|
||||
)
|
||||
logger.info("Monitor mode stopped via airmon-ng")
|
||||
except FileNotFoundError:
|
||||
try:
|
||||
subprocess.run(
|
||||
['ip', 'link', 'set', base_iface, 'down'],
|
||||
stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, timeout=10,
|
||||
)
|
||||
subprocess.run(
|
||||
['iw', 'dev', base_iface, 'set', 'type', 'managed'],
|
||||
stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, timeout=10,
|
||||
)
|
||||
subprocess.run(
|
||||
['ip', 'link', 'set', base_iface, 'up'],
|
||||
stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, timeout=10,
|
||||
)
|
||||
logger.info("Managed mode restored via iw on %s", base_iface)
|
||||
except Exception as e:
|
||||
logger.error("Failed to restore managed mode: %s", e)
|
||||
except Exception as e:
|
||||
logger.warning("airmon-ng stop failed: %s", e)
|
||||
|
||||
# Restart network services that were killed
|
||||
restarted = False
|
||||
for svc in ('wpa_supplicant', 'dhcpcd', 'NetworkManager'):
|
||||
try:
|
||||
subprocess.run(
|
||||
['systemctl', 'start', svc],
|
||||
stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, timeout=15,
|
||||
)
|
||||
restarted = True
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Wait for network services to actually reconnect before handing
|
||||
# control back so the orchestrator doesn't scan a dead interface.
|
||||
if restarted:
|
||||
logger.info("Waiting for network services to reconnect ...")
|
||||
time.sleep(5)
|
||||
|
||||
@staticmethod
|
||||
def _iface_exists(iface_name):
|
||||
"""Check if a network interface exists."""
|
||||
return os.path.isdir('/sys/class/net/%s' % iface_name)
|
||||
|
||||
# ── Bettercap subprocess management ────────────────
|
||||
|
||||
def _start_bettercap(self):
|
||||
"""Spawn bettercap subprocess with REST API."""
|
||||
cfg = self.shared_data.config
|
||||
iface = cfg.get('bifrost_iface', 'wlan0mon')
|
||||
host = cfg.get('bifrost_bettercap_host', '127.0.0.1')
|
||||
port = str(cfg.get('bifrost_bettercap_port', 8081))
|
||||
user = cfg.get('bifrost_bettercap_user', 'user')
|
||||
password = cfg.get('bifrost_bettercap_pass', 'pass')
|
||||
|
||||
cmd = [
|
||||
'bettercap', '-iface', iface, '-no-colors',
|
||||
'-eval', 'set api.rest.address %s' % host,
|
||||
'-eval', 'set api.rest.port %s' % port,
|
||||
'-eval', 'set api.rest.username %s' % user,
|
||||
'-eval', 'set api.rest.password %s' % password,
|
||||
'-eval', 'api.rest on',
|
||||
]
|
||||
|
||||
logger.info("Starting bettercap: %s", ' '.join(cmd))
|
||||
try:
|
||||
self._bettercap_proc = subprocess.Popen(
|
||||
cmd,
|
||||
stdout=subprocess.DEVNULL,
|
||||
stderr=subprocess.DEVNULL,
|
||||
)
|
||||
logger.info("bettercap PID: %d", self._bettercap_proc.pid)
|
||||
except FileNotFoundError:
|
||||
logger.error("bettercap not found! Install with: apt install bettercap")
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.error("Failed to start bettercap: %s", e)
|
||||
raise
|
||||
|
||||
def _stop_bettercap(self):
|
||||
"""Kill the bettercap subprocess."""
|
||||
if self._bettercap_proc:
|
||||
try:
|
||||
self._bettercap_proc.terminate()
|
||||
self._bettercap_proc.wait(timeout=5)
|
||||
except subprocess.TimeoutExpired:
|
||||
self._bettercap_proc.kill()
|
||||
except Exception:
|
||||
pass
|
||||
self._bettercap_proc = None
|
||||
logger.info("bettercap stopped")
|
||||
|
||||
# ── Status for web API ────────────────────────────────
|
||||
|
||||
def get_status(self):
|
||||
"""Return full engine status for web API."""
|
||||
base = {
|
||||
'enabled': self.enabled,
|
||||
'running': self._running,
|
||||
'monitor_failed': self._monitor_failed,
|
||||
}
|
||||
if self.agent and self._running:
|
||||
base.update(self.agent.get_status())
|
||||
else:
|
||||
base.update({
|
||||
'mood': 'sleeping',
|
||||
'face': '(-.-) zzZ',
|
||||
'voice': '',
|
||||
'channel': 0,
|
||||
'num_aps': 0,
|
||||
'num_handshakes': 0,
|
||||
'uptime': 0,
|
||||
'epoch': 0,
|
||||
'mode': 'auto',
|
||||
'last_pwnd': '',
|
||||
'reward': 0,
|
||||
})
|
||||
return base
|
||||
568
bifrost/agent.py
Normal file
568
bifrost/agent.py
Normal file
@@ -0,0 +1,568 @@
|
||||
"""
|
||||
Bifrost — WiFi recon agent.
|
||||
Ported from pwnagotchi/agent.py using composition instead of inheritance.
|
||||
"""
|
||||
import time
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import asyncio
|
||||
import threading
|
||||
import logging
|
||||
|
||||
from bifrost.bettercap import BettercapClient
|
||||
from bifrost.automata import BifrostAutomata
|
||||
from bifrost.epoch import BifrostEpoch
|
||||
from bifrost.voice import BifrostVoice
|
||||
from bifrost import plugins
|
||||
|
||||
from logger import Logger
|
||||
|
||||
logger = Logger(name="bifrost.agent", level=logging.DEBUG)
|
||||
|
||||
|
||||
class BifrostAgent:
|
||||
"""WiFi recon agent — drives bettercap, captures handshakes, tracks epochs."""
|
||||
|
||||
def __init__(self, shared_data, stop_event=None):
|
||||
self.shared_data = shared_data
|
||||
self._config = shared_data.config
|
||||
self.db = shared_data.db
|
||||
self._stop_event = stop_event or threading.Event()
|
||||
|
||||
# Sub-systems
|
||||
cfg = self._config
|
||||
self.bettercap = BettercapClient(
|
||||
hostname=cfg.get('bifrost_bettercap_host', '127.0.0.1'),
|
||||
scheme='http',
|
||||
port=int(cfg.get('bifrost_bettercap_port', 8081)),
|
||||
username=cfg.get('bifrost_bettercap_user', 'user'),
|
||||
password=cfg.get('bifrost_bettercap_pass', 'pass'),
|
||||
)
|
||||
self.automata = BifrostAutomata(cfg)
|
||||
self.epoch = BifrostEpoch(cfg)
|
||||
self.voice = BifrostVoice()
|
||||
|
||||
self._started_at = time.time()
|
||||
self._filter = None
|
||||
flt = cfg.get('bifrost_filter', '')
|
||||
if flt:
|
||||
try:
|
||||
self._filter = re.compile(flt)
|
||||
except re.error:
|
||||
logger.warning("Invalid bifrost_filter regex: %s", flt)
|
||||
|
||||
self._current_channel = 0
|
||||
self._tot_aps = 0
|
||||
self._aps_on_channel = 0
|
||||
self._supported_channels = list(range(1, 15))
|
||||
|
||||
self._access_points = []
|
||||
self._last_pwnd = None
|
||||
self._history = {}
|
||||
self._handshakes = {}
|
||||
self.mode = 'auto'
|
||||
|
||||
# Whitelist
|
||||
self._whitelist = [
|
||||
w.strip().lower() for w in
|
||||
str(cfg.get('bifrost_whitelist', '')).split(',') if w.strip()
|
||||
]
|
||||
# Channels
|
||||
self._channels = [
|
||||
int(c.strip()) for c in
|
||||
str(cfg.get('bifrost_channels', '')).split(',') if c.strip()
|
||||
]
|
||||
|
||||
# Ensure handshakes dir
|
||||
hs_dir = cfg.get('bifrost_bettercap_handshakes', '/root/bifrost/handshakes')
|
||||
if hs_dir and not os.path.exists(hs_dir):
|
||||
try:
|
||||
os.makedirs(hs_dir, exist_ok=True)
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
# ── Lifecycle ─────────────────────────────────────────
|
||||
|
||||
def start(self):
|
||||
"""Initialize bettercap, start monitor mode, begin event polling."""
|
||||
self._wait_bettercap()
|
||||
self.setup_events()
|
||||
self.automata.set_starting()
|
||||
self._log_activity('system', 'Bifrost starting', self.voice.on_starting())
|
||||
self.start_monitor_mode()
|
||||
self.start_event_polling()
|
||||
self.start_session_fetcher()
|
||||
self.next_epoch()
|
||||
self.automata.set_ready()
|
||||
self._log_activity('system', 'Bifrost ready', self.voice.on_ready())
|
||||
|
||||
def setup_events(self):
|
||||
"""Silence noisy bettercap events."""
|
||||
logger.info("connecting to %s ...", self.bettercap.url)
|
||||
silence = [
|
||||
'ble.device.new', 'ble.device.lost', 'ble.device.disconnected',
|
||||
'ble.device.connected', 'ble.device.service.discovered',
|
||||
'ble.device.characteristic.discovered',
|
||||
'mod.started', 'mod.stopped', 'update.available',
|
||||
'session.closing', 'session.started',
|
||||
]
|
||||
for tag in silence:
|
||||
try:
|
||||
self.bettercap.run('events.ignore %s' % tag, verbose_errors=False)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def _reset_wifi_settings(self):
|
||||
iface = self._config.get('bifrost_iface', 'wlan0mon')
|
||||
self.bettercap.run('set wifi.interface %s' % iface)
|
||||
self.bettercap.run('set wifi.ap.ttl %d' % self._config.get('bifrost_personality_ap_ttl', 120))
|
||||
self.bettercap.run('set wifi.sta.ttl %d' % self._config.get('bifrost_personality_sta_ttl', 300))
|
||||
self.bettercap.run('set wifi.rssi.min %d' % self._config.get('bifrost_personality_min_rssi', -200))
|
||||
hs_dir = self._config.get('bifrost_bettercap_handshakes', '/root/bifrost/handshakes')
|
||||
self.bettercap.run('set wifi.handshakes.file %s' % hs_dir)
|
||||
self.bettercap.run('set wifi.handshakes.aggregate false')
|
||||
|
||||
def start_monitor_mode(self):
|
||||
"""Wait for monitor interface and start wifi.recon."""
|
||||
iface = self._config.get('bifrost_iface', 'wlan0mon')
|
||||
has_mon = False
|
||||
retries = 0
|
||||
|
||||
while not has_mon and retries < 30 and not self._stop_event.is_set():
|
||||
try:
|
||||
s = self.bettercap.session()
|
||||
for i in s.get('interfaces', []):
|
||||
if i['name'] == iface:
|
||||
logger.info("found monitor interface: %s", i['name'])
|
||||
has_mon = True
|
||||
break
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
if not has_mon:
|
||||
logger.info("waiting for monitor interface %s ... (%d)", iface, retries)
|
||||
self._stop_event.wait(2)
|
||||
retries += 1
|
||||
|
||||
if not has_mon:
|
||||
logger.warning("monitor interface %s not found after %d retries", iface, retries)
|
||||
|
||||
# Detect supported channels
|
||||
try:
|
||||
from bifrost.compat import _build_utils_shim
|
||||
self._supported_channels = _build_utils_shim(self.shared_data).iface_channels(iface)
|
||||
except Exception:
|
||||
self._supported_channels = list(range(1, 15))
|
||||
|
||||
logger.info("supported channels: %s", self._supported_channels)
|
||||
self._reset_wifi_settings()
|
||||
|
||||
# Start wifi recon
|
||||
try:
|
||||
wifi_running = self._is_module_running('wifi')
|
||||
if wifi_running:
|
||||
self.bettercap.run('wifi.recon off; wifi.recon on')
|
||||
self.bettercap.run('wifi.clear')
|
||||
else:
|
||||
self.bettercap.run('wifi.recon on')
|
||||
except Exception as e:
|
||||
err_msg = str(e)
|
||||
if 'Operation not supported' in err_msg or 'EOPNOTSUPP' in err_msg:
|
||||
logger.error(
|
||||
"wifi.recon failed: %s — Your WiFi chip likely does NOT support "
|
||||
"monitor mode. The built-in Broadcom chip on Raspberry Pi Zero/Zero 2 "
|
||||
"has limited monitor mode support. Use an external USB WiFi adapter "
|
||||
"(e.g. Alfa AWUS036ACH, Panda PAU09) that supports monitor mode and "
|
||||
"packet injection.", e)
|
||||
self._log_activity('error',
|
||||
'WiFi chip does not support monitor mode',
|
||||
'Use an external USB WiFi adapter with monitor mode support')
|
||||
else:
|
||||
logger.error("Error starting wifi.recon: %s", e)
|
||||
|
||||
def _wait_bettercap(self):
|
||||
retries = 0
|
||||
while retries < 30 and not self._stop_event.is_set():
|
||||
try:
|
||||
self.bettercap.session()
|
||||
return
|
||||
except Exception:
|
||||
logger.info("waiting for bettercap API ...")
|
||||
self._stop_event.wait(2)
|
||||
retries += 1
|
||||
if not self._stop_event.is_set():
|
||||
raise Exception("bettercap API not available after 60s")
|
||||
|
||||
def _is_module_running(self, module):
|
||||
try:
|
||||
s = self.bettercap.session()
|
||||
for m in s.get('modules', []):
|
||||
if m['name'] == module:
|
||||
return m['running']
|
||||
except Exception:
|
||||
pass
|
||||
return False
|
||||
|
||||
# ── Recon cycle ───────────────────────────────────────
|
||||
|
||||
def recon(self):
|
||||
"""Full-spectrum WiFi scan for recon_time seconds."""
|
||||
recon_time = self._config.get('bifrost_personality_recon_time', 30)
|
||||
max_inactive = 3
|
||||
recon_mul = 2
|
||||
|
||||
if self.epoch.inactive_for >= max_inactive:
|
||||
recon_time *= recon_mul
|
||||
|
||||
self._current_channel = 0
|
||||
|
||||
if not self._channels:
|
||||
logger.debug("RECON %ds (all channels)", recon_time)
|
||||
try:
|
||||
self.bettercap.run('wifi.recon.channel clear')
|
||||
except Exception:
|
||||
pass
|
||||
else:
|
||||
ch_str = ','.join(map(str, self._channels))
|
||||
logger.debug("RECON %ds on channels %s", recon_time, ch_str)
|
||||
try:
|
||||
self.bettercap.run('wifi.recon.channel %s' % ch_str)
|
||||
except Exception as e:
|
||||
logger.error("Error setting recon channels: %s", e)
|
||||
|
||||
self.automata.wait_for(recon_time, self.epoch, sleeping=False,
|
||||
stop_event=self._stop_event)
|
||||
|
||||
def _filter_included(self, ap):
|
||||
if self._filter is None:
|
||||
return True
|
||||
return (self._filter.match(ap.get('hostname', '')) is not None or
|
||||
self._filter.match(ap.get('mac', '')) is not None)
|
||||
|
||||
def get_access_points(self):
|
||||
"""Fetch APs from bettercap, filter whitelist and open networks."""
|
||||
aps = []
|
||||
try:
|
||||
s = self.bettercap.session()
|
||||
plugins.on("unfiltered_ap_list", s.get('wifi', {}).get('aps', []))
|
||||
for ap in s.get('wifi', {}).get('aps', []):
|
||||
enc = ap.get('encryption', '')
|
||||
if enc == '' or enc == 'OPEN':
|
||||
continue
|
||||
hostname = ap.get('hostname', '').lower()
|
||||
mac = ap.get('mac', '').lower()
|
||||
prefix = mac[:8]
|
||||
if (hostname not in self._whitelist and
|
||||
mac not in self._whitelist and
|
||||
prefix not in self._whitelist):
|
||||
if self._filter_included(ap):
|
||||
aps.append(ap)
|
||||
except Exception as e:
|
||||
logger.error("Error getting APs: %s", e)
|
||||
|
||||
aps.sort(key=lambda a: a.get('channel', 0))
|
||||
self._access_points = aps
|
||||
plugins.on('wifi_update', aps)
|
||||
self.epoch.observe(aps, list(self.automata.peers.values()))
|
||||
|
||||
# Update DB with discovered networks
|
||||
self._persist_networks(aps)
|
||||
return aps
|
||||
|
||||
def get_access_points_by_channel(self):
|
||||
"""Get APs grouped by channel, sorted by density."""
|
||||
aps = self.get_access_points()
|
||||
grouped = {}
|
||||
for ap in aps:
|
||||
ch = ap.get('channel', 0)
|
||||
if self._channels and ch not in self._channels:
|
||||
continue
|
||||
grouped.setdefault(ch, []).append(ap)
|
||||
return sorted(grouped.items(), key=lambda kv: len(kv[1]), reverse=True)
|
||||
|
||||
# ── Actions ───────────────────────────────────────────
|
||||
|
||||
def _should_interact(self, who):
|
||||
if self._has_handshake(who):
|
||||
return False
|
||||
if who not in self._history:
|
||||
self._history[who] = 1
|
||||
return True
|
||||
self._history[who] += 1
|
||||
max_int = self._config.get('bifrost_personality_max_interactions', 3)
|
||||
return self._history[who] < max_int
|
||||
|
||||
def _has_handshake(self, bssid):
|
||||
for key in self._handshakes:
|
||||
if bssid.lower() in key:
|
||||
return True
|
||||
return False
|
||||
|
||||
def associate(self, ap, throttle=0):
|
||||
"""Send association frame to trigger PMKID."""
|
||||
if self.automata.is_stale(self.epoch):
|
||||
return
|
||||
if (self._config.get('bifrost_personality_associate', True) and
|
||||
self._should_interact(ap.get('mac', ''))):
|
||||
try:
|
||||
hostname = ap.get('hostname', ap.get('mac', '?'))
|
||||
logger.info("ASSOC %s (%s) ch=%d rssi=%d",
|
||||
hostname, ap.get('mac', ''), ap.get('channel', 0), ap.get('rssi', 0))
|
||||
self.bettercap.run('wifi.assoc %s' % ap['mac'])
|
||||
self.epoch.track(assoc=True)
|
||||
self._log_activity('assoc', 'Association: %s' % hostname,
|
||||
self.voice.on_assoc(hostname))
|
||||
except Exception as e:
|
||||
self.automata.on_error(ap.get('mac', ''), e)
|
||||
plugins.on('association', ap)
|
||||
if throttle > 0:
|
||||
time.sleep(throttle)
|
||||
|
||||
def deauth(self, ap, sta, throttle=0):
|
||||
"""Deauthenticate client to capture handshake."""
|
||||
if self.automata.is_stale(self.epoch):
|
||||
return
|
||||
if (self._config.get('bifrost_personality_deauth', True) and
|
||||
self._should_interact(sta.get('mac', ''))):
|
||||
try:
|
||||
logger.info("DEAUTH %s (%s) from %s ch=%d",
|
||||
sta.get('mac', ''), sta.get('vendor', ''),
|
||||
ap.get('hostname', ap.get('mac', '')), ap.get('channel', 0))
|
||||
self.bettercap.run('wifi.deauth %s' % sta['mac'])
|
||||
self.epoch.track(deauth=True)
|
||||
self._log_activity('deauth', 'Deauth: %s' % sta.get('mac', ''),
|
||||
self.voice.on_deauth(sta.get('mac', '')))
|
||||
except Exception as e:
|
||||
self.automata.on_error(sta.get('mac', ''), e)
|
||||
plugins.on('deauthentication', ap, sta)
|
||||
if throttle > 0:
|
||||
time.sleep(throttle)
|
||||
|
||||
def set_channel(self, channel, verbose=True):
|
||||
"""Hop to a specific WiFi channel."""
|
||||
if self.automata.is_stale(self.epoch):
|
||||
return
|
||||
wait = 0
|
||||
if self.epoch.did_deauth:
|
||||
wait = self._config.get('bifrost_personality_hop_recon_time', 10)
|
||||
elif self.epoch.did_associate:
|
||||
wait = self._config.get('bifrost_personality_min_recon_time', 5)
|
||||
|
||||
if channel != self._current_channel:
|
||||
if self._current_channel != 0 and wait > 0:
|
||||
logger.debug("waiting %ds on channel %d", wait, self._current_channel)
|
||||
self.automata.wait_for(wait, self.epoch, stop_event=self._stop_event)
|
||||
try:
|
||||
self.bettercap.run('wifi.recon.channel %d' % channel)
|
||||
self._current_channel = channel
|
||||
self.epoch.track(hop=True)
|
||||
plugins.on('channel_hop', channel)
|
||||
except Exception as e:
|
||||
logger.error("Error setting channel: %s", e)
|
||||
|
||||
def next_epoch(self):
|
||||
"""Transition to next epoch — evaluate mood."""
|
||||
self.automata.next_epoch(self.epoch)
|
||||
# Persist epoch to DB
|
||||
data = self.epoch.data()
|
||||
self._persist_epoch(data)
|
||||
self._log_activity('epoch', 'Epoch %d' % (self.epoch.epoch - 1),
|
||||
self.voice.on_epoch(self.epoch.epoch - 1))
|
||||
|
||||
# ── Event polling ─────────────────────────────────────
|
||||
|
||||
def start_event_polling(self):
|
||||
"""Start event listener in background thread.
|
||||
|
||||
Tries websocket first; falls back to REST polling if the
|
||||
``websockets`` package is not installed.
|
||||
"""
|
||||
t = threading.Thread(target=self._event_poller, daemon=True, name="BifrostEvents")
|
||||
t.start()
|
||||
|
||||
def _event_poller(self):
|
||||
try:
|
||||
self.bettercap.run('events.clear')
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Probe once whether websockets is available
|
||||
try:
|
||||
import websockets # noqa: F401
|
||||
has_ws = True
|
||||
except ImportError:
|
||||
has_ws = False
|
||||
logger.warning("websockets package not installed — using REST event polling "
|
||||
"(pip install websockets for real-time events)")
|
||||
|
||||
if has_ws:
|
||||
self._ws_event_loop()
|
||||
else:
|
||||
self._rest_event_loop()
|
||||
|
||||
def _ws_event_loop(self):
|
||||
"""Websocket-based event listener (preferred)."""
|
||||
loop = asyncio.new_event_loop()
|
||||
asyncio.set_event_loop(loop)
|
||||
while not self._stop_event.is_set():
|
||||
try:
|
||||
loop.run_until_complete(self.bettercap.start_websocket(
|
||||
self._on_event, self._stop_event))
|
||||
except Exception as ex:
|
||||
if self._stop_event.is_set():
|
||||
break
|
||||
logger.debug("Event poller error: %s", ex)
|
||||
self._stop_event.wait(5)
|
||||
loop.close()
|
||||
|
||||
def _rest_event_loop(self):
|
||||
"""REST-based fallback event poller — polls /api/events every 2s."""
|
||||
while not self._stop_event.is_set():
|
||||
try:
|
||||
events = self.bettercap.events()
|
||||
for ev in (events or []):
|
||||
tag = ev.get('tag', '')
|
||||
if tag == 'wifi.client.handshake':
|
||||
# Build a fake websocket message for the existing handler
|
||||
import asyncio as _aio
|
||||
_loop = _aio.new_event_loop()
|
||||
_loop.run_until_complete(self._on_event(json.dumps(ev)))
|
||||
_loop.close()
|
||||
except Exception as ex:
|
||||
logger.debug("REST event poll error: %s", ex)
|
||||
self._stop_event.wait(2)
|
||||
|
||||
async def _on_event(self, msg):
|
||||
"""Handle bettercap websocket events."""
|
||||
try:
|
||||
jmsg = json.loads(msg)
|
||||
except json.JSONDecodeError:
|
||||
return
|
||||
|
||||
if jmsg.get('tag') == 'wifi.client.handshake':
|
||||
filename = jmsg.get('data', {}).get('file', '')
|
||||
sta_mac = jmsg.get('data', {}).get('station', '')
|
||||
ap_mac = jmsg.get('data', {}).get('ap', '')
|
||||
key = "%s -> %s" % (sta_mac, ap_mac)
|
||||
|
||||
if key not in self._handshakes:
|
||||
self._handshakes[key] = jmsg
|
||||
self._last_pwnd = ap_mac
|
||||
|
||||
# Find AP info
|
||||
ap_name = ap_mac
|
||||
try:
|
||||
s = self.bettercap.session()
|
||||
for ap in s.get('wifi', {}).get('aps', []):
|
||||
if ap.get('mac') == ap_mac:
|
||||
if ap.get('hostname') and ap['hostname'] != '<hidden>':
|
||||
ap_name = ap['hostname']
|
||||
break
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
logger.warning("!!! HANDSHAKE: %s -> %s !!!", sta_mac, ap_name)
|
||||
self.epoch.track(handshake=True)
|
||||
self._persist_handshake(ap_mac, sta_mac, ap_name, filename)
|
||||
self._log_activity('handshake',
|
||||
'Handshake: %s' % ap_name,
|
||||
self.voice.on_handshakes(1))
|
||||
plugins.on('handshake', filename, ap_mac, sta_mac)
|
||||
|
||||
def start_session_fetcher(self):
|
||||
"""Start background thread that polls bettercap for stats."""
|
||||
t = threading.Thread(target=self._fetch_stats, daemon=True, name="BifrostStats")
|
||||
t.start()
|
||||
|
||||
def _fetch_stats(self):
|
||||
while not self._stop_event.is_set():
|
||||
try:
|
||||
s = self.bettercap.session()
|
||||
self._tot_aps = len(s.get('wifi', {}).get('aps', []))
|
||||
except Exception:
|
||||
pass
|
||||
self._stop_event.wait(2)
|
||||
|
||||
# ── Status for web API ────────────────────────────────
|
||||
|
||||
def get_status(self):
|
||||
"""Return current agent state for the web API."""
|
||||
return {
|
||||
'mood': self.automata.mood,
|
||||
'face': self.automata.face,
|
||||
'voice': self.automata.voice_text,
|
||||
'channel': self._current_channel,
|
||||
'num_aps': self._tot_aps,
|
||||
'num_handshakes': len(self._handshakes),
|
||||
'uptime': int(time.time() - self._started_at),
|
||||
'epoch': self.epoch.epoch,
|
||||
'mode': self.mode,
|
||||
'last_pwnd': self._last_pwnd or '',
|
||||
'reward': self.epoch.data().get('reward', 0),
|
||||
}
|
||||
|
||||
# ── DB persistence ────────────────────────────────────
|
||||
|
||||
def _persist_networks(self, aps):
|
||||
"""Upsert discovered networks to DB."""
|
||||
for ap in aps:
|
||||
try:
|
||||
self.db.execute(
|
||||
"""INSERT INTO bifrost_networks
|
||||
(bssid, essid, channel, encryption, rssi, vendor, num_clients, last_seen)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, CURRENT_TIMESTAMP)
|
||||
ON CONFLICT(bssid) DO UPDATE SET
|
||||
essid=?, channel=?, encryption=?, rssi=?, vendor=?,
|
||||
num_clients=?, last_seen=CURRENT_TIMESTAMP""",
|
||||
(ap.get('mac', ''), ap.get('hostname', ''), ap.get('channel', 0),
|
||||
ap.get('encryption', ''), ap.get('rssi', 0), ap.get('vendor', ''),
|
||||
len(ap.get('clients', [])),
|
||||
ap.get('hostname', ''), ap.get('channel', 0),
|
||||
ap.get('encryption', ''), ap.get('rssi', 0), ap.get('vendor', ''),
|
||||
len(ap.get('clients', [])))
|
||||
)
|
||||
except Exception as e:
|
||||
logger.debug("Error persisting network: %s", e)
|
||||
|
||||
def _persist_handshake(self, ap_mac, sta_mac, ap_name, filename):
|
||||
try:
|
||||
self.db.execute(
|
||||
"""INSERT OR IGNORE INTO bifrost_handshakes
|
||||
(ap_mac, sta_mac, ap_essid, filename)
|
||||
VALUES (?, ?, ?, ?)""",
|
||||
(ap_mac, sta_mac, ap_name, filename)
|
||||
)
|
||||
except Exception as e:
|
||||
logger.debug("Error persisting handshake: %s", e)
|
||||
|
||||
def _persist_epoch(self, data):
|
||||
try:
|
||||
self.db.execute(
|
||||
"""INSERT INTO bifrost_epochs
|
||||
(epoch_num, started_at, duration_secs, num_deauths, num_assocs,
|
||||
num_handshakes, num_hops, num_missed, num_peers, mood, reward,
|
||||
cpu_load, mem_usage, temperature, meta_json)
|
||||
VALUES (?, datetime('now'), ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""",
|
||||
(self.epoch.epoch - 1, data.get('duration_secs', 0),
|
||||
data.get('num_deauths', 0), data.get('num_associations', 0),
|
||||
data.get('num_handshakes', 0), data.get('num_hops', 0),
|
||||
data.get('missed_interactions', 0), data.get('num_peers', 0),
|
||||
self.automata.mood, data.get('reward', 0),
|
||||
data.get('cpu_load', 0), data.get('mem_usage', 0),
|
||||
data.get('temperature', 0), '{}')
|
||||
)
|
||||
except Exception as e:
|
||||
logger.debug("Error persisting epoch: %s", e)
|
||||
|
||||
def _log_activity(self, event_type, title, details=''):
|
||||
"""Log an activity event to the DB."""
|
||||
self.automata.voice_text = details or title
|
||||
try:
|
||||
self.db.execute(
|
||||
"""INSERT INTO bifrost_activity (event_type, title, details)
|
||||
VALUES (?, ?, ?)""",
|
||||
(event_type, title, details)
|
||||
)
|
||||
except Exception as e:
|
||||
logger.debug("Error logging activity: %s", e)
|
||||
168
bifrost/automata.py
Normal file
168
bifrost/automata.py
Normal file
@@ -0,0 +1,168 @@
|
||||
"""
|
||||
Bifrost — Mood state machine.
|
||||
Ported from pwnagotchi/automata.py.
|
||||
"""
|
||||
import logging
|
||||
|
||||
from bifrost import plugins as plugins
|
||||
from bifrost.faces import MOOD_FACES
|
||||
|
||||
from logger import Logger
|
||||
|
||||
logger = Logger(name="bifrost.automata", level=logging.DEBUG)
|
||||
|
||||
|
||||
class BifrostAutomata:
|
||||
"""Evaluates epoch data and transitions between moods."""
|
||||
|
||||
def __init__(self, config):
|
||||
self._config = config
|
||||
self.mood = 'starting'
|
||||
self.face = MOOD_FACES.get('starting', '(. .)')
|
||||
self.voice_text = ''
|
||||
self._peers = {} # peer_id -> peer_data
|
||||
|
||||
@property
|
||||
def peers(self):
|
||||
return self._peers
|
||||
|
||||
def _set_mood(self, mood):
|
||||
self.mood = mood
|
||||
self.face = MOOD_FACES.get(mood, '(. .)')
|
||||
|
||||
def set_starting(self):
|
||||
self._set_mood('starting')
|
||||
|
||||
def set_ready(self):
|
||||
self._set_mood('ready')
|
||||
plugins.on('ready')
|
||||
|
||||
def _has_support_network_for(self, factor):
|
||||
bond_factor = self._config.get('bifrost_personality_bond_factor', 20000)
|
||||
total_encounters = sum(
|
||||
p.get('encounters', 0) if isinstance(p, dict) else getattr(p, 'encounters', 0)
|
||||
for p in self._peers.values()
|
||||
)
|
||||
support_factor = total_encounters / bond_factor
|
||||
return support_factor >= factor
|
||||
|
||||
def in_good_mood(self):
|
||||
return self._has_support_network_for(1.0)
|
||||
|
||||
def set_grateful(self):
|
||||
self._set_mood('grateful')
|
||||
plugins.on('grateful')
|
||||
|
||||
def set_lonely(self):
|
||||
if not self._has_support_network_for(1.0):
|
||||
logger.info("unit is lonely")
|
||||
self._set_mood('lonely')
|
||||
plugins.on('lonely')
|
||||
else:
|
||||
logger.info("unit is grateful instead of lonely")
|
||||
self.set_grateful()
|
||||
|
||||
def set_bored(self, inactive_for):
|
||||
bored_epochs = self._config.get('bifrost_personality_bored_epochs', 15)
|
||||
factor = inactive_for / bored_epochs if bored_epochs else 1
|
||||
if not self._has_support_network_for(factor):
|
||||
logger.warning("%d epochs with no activity -> bored", inactive_for)
|
||||
self._set_mood('bored')
|
||||
plugins.on('bored')
|
||||
else:
|
||||
logger.info("unit is grateful instead of bored")
|
||||
self.set_grateful()
|
||||
|
||||
def set_sad(self, inactive_for):
|
||||
sad_epochs = self._config.get('bifrost_personality_sad_epochs', 25)
|
||||
factor = inactive_for / sad_epochs if sad_epochs else 1
|
||||
if not self._has_support_network_for(factor):
|
||||
logger.warning("%d epochs with no activity -> sad", inactive_for)
|
||||
self._set_mood('sad')
|
||||
plugins.on('sad')
|
||||
else:
|
||||
logger.info("unit is grateful instead of sad")
|
||||
self.set_grateful()
|
||||
|
||||
def set_angry(self, factor):
|
||||
if not self._has_support_network_for(factor):
|
||||
logger.warning("too many misses -> angry (factor=%.1f)", factor)
|
||||
self._set_mood('angry')
|
||||
plugins.on('angry')
|
||||
else:
|
||||
logger.info("unit is grateful instead of angry")
|
||||
self.set_grateful()
|
||||
|
||||
def set_excited(self):
|
||||
logger.warning("lots of activity -> excited")
|
||||
self._set_mood('excited')
|
||||
plugins.on('excited')
|
||||
|
||||
def set_rebooting(self):
|
||||
self._set_mood('broken')
|
||||
plugins.on('rebooting')
|
||||
|
||||
def next_epoch(self, epoch):
|
||||
"""Evaluate epoch state and transition mood.
|
||||
|
||||
Args:
|
||||
epoch: BifrostEpoch instance
|
||||
"""
|
||||
was_stale = epoch.num_missed > self._config.get('bifrost_personality_max_misses', 8)
|
||||
did_miss = epoch.num_missed
|
||||
|
||||
# Trigger epoch transition (resets counters, computes reward)
|
||||
epoch.next()
|
||||
|
||||
max_misses = self._config.get('bifrost_personality_max_misses', 8)
|
||||
excited_threshold = self._config.get('bifrost_personality_excited_epochs', 10)
|
||||
|
||||
# Mood evaluation (same logic as pwnagotchi automata.py)
|
||||
if was_stale:
|
||||
factor = did_miss / max_misses if max_misses else 1
|
||||
if factor >= 2.0:
|
||||
self.set_angry(factor)
|
||||
else:
|
||||
logger.warning("agent missed %d interactions -> lonely", did_miss)
|
||||
self.set_lonely()
|
||||
elif epoch.sad_for:
|
||||
sad_epochs = self._config.get('bifrost_personality_sad_epochs', 25)
|
||||
factor = epoch.inactive_for / sad_epochs if sad_epochs else 1
|
||||
if factor >= 2.0:
|
||||
self.set_angry(factor)
|
||||
else:
|
||||
self.set_sad(epoch.inactive_for)
|
||||
elif epoch.bored_for:
|
||||
self.set_bored(epoch.inactive_for)
|
||||
elif epoch.active_for >= excited_threshold:
|
||||
self.set_excited()
|
||||
elif epoch.active_for >= 5 and self._has_support_network_for(5.0):
|
||||
self.set_grateful()
|
||||
|
||||
plugins.on('epoch', epoch.epoch - 1, epoch.data())
|
||||
|
||||
def on_miss(self, who):
|
||||
logger.info("it looks like %s is not in range anymore :/", who)
|
||||
|
||||
def on_error(self, who, e):
|
||||
if 'is an unknown BSSID' in str(e):
|
||||
self.on_miss(who)
|
||||
else:
|
||||
logger.error(str(e))
|
||||
|
||||
def is_stale(self, epoch):
|
||||
return epoch.num_missed > self._config.get('bifrost_personality_max_misses', 8)
|
||||
|
||||
def wait_for(self, t, epoch, sleeping=True, stop_event=None):
|
||||
"""Wait and track sleep time.
|
||||
|
||||
If *stop_event* is provided the wait is interruptible so the
|
||||
engine can shut down quickly even during long recon windows.
|
||||
"""
|
||||
plugins.on('sleep' if sleeping else 'wait', t)
|
||||
epoch.track(sleep=True, inc=t)
|
||||
import time
|
||||
if stop_event is not None:
|
||||
stop_event.wait(t)
|
||||
else:
|
||||
time.sleep(t)
|
||||
103
bifrost/bettercap.py
Normal file
103
bifrost/bettercap.py
Normal file
@@ -0,0 +1,103 @@
|
||||
"""
|
||||
Bifrost — Bettercap REST API client.
|
||||
Ported from pwnagotchi/bettercap.py using urllib (no requests dependency).
|
||||
"""
|
||||
import json
|
||||
import logging
|
||||
import base64
|
||||
import urllib.request
|
||||
import urllib.error
|
||||
|
||||
from logger import Logger
|
||||
|
||||
logger = Logger(name="bifrost.bettercap", level=logging.DEBUG)
|
||||
|
||||
|
||||
class BettercapClient:
|
||||
"""Synchronous REST client for the bettercap API."""
|
||||
|
||||
def __init__(self, hostname='127.0.0.1', scheme='http', port=8081,
|
||||
username='user', password='pass'):
|
||||
self.hostname = hostname
|
||||
self.scheme = scheme
|
||||
self.port = port
|
||||
self.username = username
|
||||
self.password = password
|
||||
self.url = "%s://%s:%d/api" % (scheme, hostname, port)
|
||||
self.websocket = "ws://%s:%s@%s:%d/api" % (username, password, hostname, port)
|
||||
self._auth_header = 'Basic ' + base64.b64encode(
|
||||
('%s:%s' % (username, password)).encode()
|
||||
).decode()
|
||||
|
||||
def _request(self, method, path, data=None, verbose_errors=True):
|
||||
"""Make an HTTP request to bettercap API."""
|
||||
url = "%s%s" % (self.url, path)
|
||||
body = json.dumps(data).encode() if data else None
|
||||
req = urllib.request.Request(url, data=body, method=method)
|
||||
req.add_header('Authorization', self._auth_header)
|
||||
if body:
|
||||
req.add_header('Content-Type', 'application/json')
|
||||
|
||||
try:
|
||||
with urllib.request.urlopen(req, timeout=10) as resp:
|
||||
raw = resp.read().decode('utf-8')
|
||||
try:
|
||||
return json.loads(raw)
|
||||
except json.JSONDecodeError:
|
||||
return raw
|
||||
except urllib.error.HTTPError as e:
|
||||
err = "error %d: %s" % (e.code, e.read().decode('utf-8', errors='replace').strip())
|
||||
if verbose_errors:
|
||||
logger.info(err)
|
||||
raise Exception(err)
|
||||
except urllib.error.URLError as e:
|
||||
raise Exception("bettercap unreachable: %s" % e.reason)
|
||||
|
||||
def session(self):
|
||||
"""GET /api/session — current bettercap state."""
|
||||
return self._request('GET', '/session')
|
||||
|
||||
def run(self, command, verbose_errors=True):
|
||||
"""POST /api/session — execute a bettercap command."""
|
||||
return self._request('POST', '/session', {'cmd': command},
|
||||
verbose_errors=verbose_errors)
|
||||
|
||||
def events(self):
|
||||
"""GET /api/events — poll recent events (REST fallback)."""
|
||||
try:
|
||||
result = self._request('GET', '/events', verbose_errors=False)
|
||||
# Clear after reading so we don't reprocess
|
||||
try:
|
||||
self.run('events.clear', verbose_errors=False)
|
||||
except Exception:
|
||||
pass
|
||||
return result if isinstance(result, list) else []
|
||||
except Exception:
|
||||
return []
|
||||
|
||||
async def start_websocket(self, consumer, stop_event=None):
|
||||
"""Connect to bettercap websocket event stream.
|
||||
|
||||
Args:
|
||||
consumer: async callable that receives each message string.
|
||||
stop_event: optional threading.Event — exit when set.
|
||||
"""
|
||||
import websockets
|
||||
import asyncio
|
||||
ws_url = "%s/events" % self.websocket
|
||||
while not (stop_event and stop_event.is_set()):
|
||||
try:
|
||||
async with websockets.connect(ws_url, ping_interval=60,
|
||||
ping_timeout=90) as ws:
|
||||
async for msg in ws:
|
||||
if stop_event and stop_event.is_set():
|
||||
return
|
||||
try:
|
||||
await consumer(msg)
|
||||
except Exception as ex:
|
||||
logger.debug("Error parsing event: %s", ex)
|
||||
except Exception as ex:
|
||||
if stop_event and stop_event.is_set():
|
||||
return
|
||||
logger.debug("Websocket error: %s — reconnecting...", ex)
|
||||
await asyncio.sleep(2)
|
||||
185
bifrost/compat.py
Normal file
185
bifrost/compat.py
Normal file
@@ -0,0 +1,185 @@
|
||||
"""
|
||||
Bifrost — Pwnagotchi compatibility shim.
|
||||
Registers `pwnagotchi` in sys.modules so existing plugins can
|
||||
`import pwnagotchi` and get Bifrost-backed implementations.
|
||||
"""
|
||||
import sys
|
||||
import time
|
||||
import types
|
||||
import os
|
||||
|
||||
|
||||
def install_shim(shared_data, bifrost_plugins_module):
|
||||
"""Install the pwnagotchi namespace shim into sys.modules.
|
||||
|
||||
Call this BEFORE loading any pwnagotchi plugins so their
|
||||
`import pwnagotchi` resolves to our shim.
|
||||
"""
|
||||
_start_time = time.time()
|
||||
|
||||
# Create the fake pwnagotchi module
|
||||
pwn = types.ModuleType('pwnagotchi')
|
||||
pwn.__version__ = '2.0.0-bifrost'
|
||||
pwn.__file__ = __file__
|
||||
pwn.config = _build_compat_config(shared_data)
|
||||
|
||||
def _name():
|
||||
return shared_data.config.get('bjorn_name', 'bifrost')
|
||||
|
||||
def _set_name(n):
|
||||
pass # no-op, name comes from Bjorn config
|
||||
|
||||
def _uptime():
|
||||
return time.time() - _start_time
|
||||
|
||||
def _cpu_load():
|
||||
try:
|
||||
return os.getloadavg()[0]
|
||||
except (OSError, AttributeError):
|
||||
return 0.0
|
||||
|
||||
def _mem_usage():
|
||||
try:
|
||||
with open('/proc/meminfo', 'r') as f:
|
||||
lines = f.readlines()
|
||||
total = int(lines[0].split()[1])
|
||||
available = int(lines[2].split()[1])
|
||||
return (total - available) / total if total else 0.0
|
||||
except Exception:
|
||||
return 0.0
|
||||
|
||||
def _temperature():
|
||||
try:
|
||||
with open('/sys/class/thermal/thermal_zone0/temp', 'r') as f:
|
||||
return int(f.read().strip()) / 1000.0
|
||||
except Exception:
|
||||
return 0.0
|
||||
|
||||
def _reboot():
|
||||
pass # no-op in Bifrost — we don't auto-reboot
|
||||
|
||||
pwn.name = _name
|
||||
pwn.set_name = _set_name
|
||||
pwn.uptime = _uptime
|
||||
pwn.cpu_load = _cpu_load
|
||||
pwn.mem_usage = _mem_usage
|
||||
pwn.temperature = _temperature
|
||||
pwn.reboot = _reboot
|
||||
|
||||
# Register modules
|
||||
sys.modules['pwnagotchi'] = pwn
|
||||
sys.modules['pwnagotchi.plugins'] = bifrost_plugins_module
|
||||
sys.modules['pwnagotchi.utils'] = _build_utils_shim(shared_data)
|
||||
|
||||
|
||||
def _build_compat_config(shared_data):
|
||||
"""Translate Bjorn's flat bifrost_* config to pwnagotchi's nested format."""
|
||||
cfg = shared_data.config
|
||||
return {
|
||||
'main': {
|
||||
'name': cfg.get('bjorn_name', 'bifrost'),
|
||||
'iface': cfg.get('bifrost_iface', 'wlan0mon'),
|
||||
'mon_start_cmd': '',
|
||||
'no_restart': False,
|
||||
'filter': cfg.get('bifrost_filter', ''),
|
||||
'whitelist': [
|
||||
w.strip() for w in
|
||||
str(cfg.get('bifrost_whitelist', '')).split(',') if w.strip()
|
||||
],
|
||||
'plugins': cfg.get('bifrost_plugins', {}),
|
||||
'custom_plugins': cfg.get('bifrost_plugins_path', ''),
|
||||
'mon_max_blind_epochs': 50,
|
||||
},
|
||||
'personality': {
|
||||
'ap_ttl': cfg.get('bifrost_personality_ap_ttl', 120),
|
||||
'sta_ttl': cfg.get('bifrost_personality_sta_ttl', 300),
|
||||
'min_rssi': cfg.get('bifrost_personality_min_rssi', -200),
|
||||
'associate': cfg.get('bifrost_personality_associate', True),
|
||||
'deauth': cfg.get('bifrost_personality_deauth', True),
|
||||
'recon_time': cfg.get('bifrost_personality_recon_time', 30),
|
||||
'hop_recon_time': cfg.get('bifrost_personality_hop_recon_time', 10),
|
||||
'min_recon_time': cfg.get('bifrost_personality_min_recon_time', 5),
|
||||
'max_inactive_scale': 3,
|
||||
'recon_inactive_multiplier': 2,
|
||||
'max_interactions': cfg.get('bifrost_personality_max_interactions', 3),
|
||||
'max_misses_for_recon': cfg.get('bifrost_personality_max_misses', 8),
|
||||
'excited_num_epochs': cfg.get('bifrost_personality_excited_epochs', 10),
|
||||
'bored_num_epochs': cfg.get('bifrost_personality_bored_epochs', 15),
|
||||
'sad_num_epochs': cfg.get('bifrost_personality_sad_epochs', 25),
|
||||
'bond_encounters_factor': cfg.get('bifrost_personality_bond_factor', 20000),
|
||||
'channels': [
|
||||
int(c.strip()) for c in
|
||||
str(cfg.get('bifrost_channels', '')).split(',') if c.strip()
|
||||
],
|
||||
},
|
||||
'bettercap': {
|
||||
'hostname': cfg.get('bifrost_bettercap_host', '127.0.0.1'),
|
||||
'scheme': 'http',
|
||||
'port': cfg.get('bifrost_bettercap_port', 8081),
|
||||
'username': cfg.get('bifrost_bettercap_user', 'user'),
|
||||
'password': cfg.get('bifrost_bettercap_pass', 'pass'),
|
||||
'handshakes': cfg.get('bifrost_bettercap_handshakes', '/root/bifrost/handshakes'),
|
||||
'silence': [
|
||||
'ble.device.new', 'ble.device.lost', 'ble.device.disconnected',
|
||||
'ble.device.connected', 'ble.device.service.discovered',
|
||||
'ble.device.characteristic.discovered',
|
||||
'mod.started', 'mod.stopped', 'update.available',
|
||||
'session.closing', 'session.started',
|
||||
],
|
||||
},
|
||||
'ai': {
|
||||
'enabled': cfg.get('bifrost_ai_enabled', False),
|
||||
'path': '/root/bifrost/brain.json',
|
||||
},
|
||||
'ui': {
|
||||
'fps': 1.0,
|
||||
'web': {'enabled': False},
|
||||
'display': {'enabled': False},
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def _build_utils_shim(shared_data):
|
||||
"""Minimal pwnagotchi.utils shim."""
|
||||
mod = types.ModuleType('pwnagotchi.utils')
|
||||
|
||||
def secs_to_hhmmss(secs):
|
||||
h = int(secs // 3600)
|
||||
m = int((secs % 3600) // 60)
|
||||
s = int(secs % 60)
|
||||
return "%d:%02d:%02d" % (h, m, s)
|
||||
|
||||
def iface_channels(iface):
|
||||
"""Return available channels for interface."""
|
||||
try:
|
||||
import subprocess
|
||||
out = subprocess.check_output(
|
||||
['iwlist', iface, 'channel'],
|
||||
stderr=subprocess.DEVNULL, timeout=5
|
||||
).decode()
|
||||
channels = []
|
||||
for line in out.split('\n'):
|
||||
if 'Channel' in line and 'Current' not in line:
|
||||
parts = line.strip().split()
|
||||
for p in parts:
|
||||
try:
|
||||
ch = int(p)
|
||||
if 1 <= ch <= 14:
|
||||
channels.append(ch)
|
||||
except ValueError:
|
||||
continue
|
||||
return sorted(set(channels)) if channels else list(range(1, 15))
|
||||
except Exception:
|
||||
return list(range(1, 15))
|
||||
|
||||
def total_unique_handshakes(path):
|
||||
"""Count unique handshake files in directory."""
|
||||
import glob as _glob
|
||||
if not os.path.isdir(path):
|
||||
return 0
|
||||
return len(_glob.glob(os.path.join(path, '*.pcap')))
|
||||
|
||||
mod.secs_to_hhmmss = secs_to_hhmmss
|
||||
mod.iface_channels = iface_channels
|
||||
mod.total_unique_handshakes = total_unique_handshakes
|
||||
return mod
|
||||
292
bifrost/epoch.py
Normal file
292
bifrost/epoch.py
Normal file
@@ -0,0 +1,292 @@
|
||||
"""
|
||||
Bifrost — Epoch tracking.
|
||||
Ported from pwnagotchi/ai/epoch.py + pwnagotchi/ai/reward.py.
|
||||
"""
|
||||
import time
|
||||
import threading
|
||||
import logging
|
||||
import os
|
||||
|
||||
from logger import Logger
|
||||
|
||||
logger = Logger(name="bifrost.epoch", level=logging.DEBUG)
|
||||
|
||||
NUM_CHANNELS = 14 # 2.4 GHz channels
|
||||
|
||||
|
||||
# ── Reward function (from pwnagotchi/ai/reward.py) ──────────────
|
||||
|
||||
class RewardFunction:
|
||||
"""Reward signal for RL — higher is better."""
|
||||
|
||||
def __call__(self, epoch_n, state):
|
||||
eps = 1e-20
|
||||
tot_epochs = epoch_n + eps
|
||||
tot_interactions = max(
|
||||
state['num_deauths'] + state['num_associations'],
|
||||
state['num_handshakes']
|
||||
) + eps
|
||||
tot_channels = NUM_CHANNELS
|
||||
|
||||
# Positive signals
|
||||
h = state['num_handshakes'] / tot_interactions
|
||||
a = 0.2 * (state['active_for_epochs'] / tot_epochs)
|
||||
c = 0.1 * (state['num_hops'] / tot_channels)
|
||||
|
||||
# Negative signals
|
||||
b = -0.3 * (state['blind_for_epochs'] / tot_epochs)
|
||||
m = -0.3 * (state['missed_interactions'] / tot_interactions)
|
||||
i = -0.2 * (state['inactive_for_epochs'] / tot_epochs)
|
||||
|
||||
_sad = state['sad_for_epochs'] if state['sad_for_epochs'] >= 5 else 0
|
||||
_bored = state['bored_for_epochs'] if state['bored_for_epochs'] >= 5 else 0
|
||||
s = -0.2 * (_sad / tot_epochs)
|
||||
l_val = -0.1 * (_bored / tot_epochs)
|
||||
|
||||
return h + a + c + b + i + m + s + l_val
|
||||
|
||||
|
||||
# ── Epoch state ──────────────────────────────────────────────────
|
||||
|
||||
class BifrostEpoch:
|
||||
"""Tracks per-epoch counters, observations, and reward."""
|
||||
|
||||
def __init__(self, config):
|
||||
self.epoch = 0
|
||||
self.config = config
|
||||
|
||||
# Consecutive epoch counters
|
||||
self.inactive_for = 0
|
||||
self.active_for = 0
|
||||
self.blind_for = 0
|
||||
self.sad_for = 0
|
||||
self.bored_for = 0
|
||||
|
||||
# Per-epoch action flags & counters
|
||||
self.did_deauth = False
|
||||
self.num_deauths = 0
|
||||
self.did_associate = False
|
||||
self.num_assocs = 0
|
||||
self.num_missed = 0
|
||||
self.did_handshakes = False
|
||||
self.num_shakes = 0
|
||||
self.num_hops = 0
|
||||
self.num_slept = 0
|
||||
self.num_peers = 0
|
||||
self.tot_bond_factor = 0.0
|
||||
self.avg_bond_factor = 0.0
|
||||
self.any_activity = False
|
||||
|
||||
# Timing
|
||||
self.epoch_started = time.time()
|
||||
self.epoch_duration = 0
|
||||
|
||||
# Channel histograms for AI observation
|
||||
self.non_overlapping_channels = {1: 0, 6: 0, 11: 0}
|
||||
self._observation = {
|
||||
'aps_histogram': [0.0] * NUM_CHANNELS,
|
||||
'sta_histogram': [0.0] * NUM_CHANNELS,
|
||||
'peers_histogram': [0.0] * NUM_CHANNELS,
|
||||
}
|
||||
self._observation_ready = threading.Event()
|
||||
self._epoch_data = {}
|
||||
self._epoch_data_ready = threading.Event()
|
||||
self._reward = RewardFunction()
|
||||
|
||||
def wait_for_epoch_data(self, with_observation=True, timeout=None):
|
||||
self._epoch_data_ready.wait(timeout)
|
||||
self._epoch_data_ready.clear()
|
||||
if with_observation:
|
||||
return {**self._observation, **self._epoch_data}
|
||||
return self._epoch_data
|
||||
|
||||
def data(self):
|
||||
return self._epoch_data
|
||||
|
||||
def observe(self, aps, peers):
|
||||
"""Update observation histograms from current AP/peer lists."""
|
||||
num_aps = len(aps)
|
||||
if num_aps == 0:
|
||||
self.blind_for += 1
|
||||
else:
|
||||
self.blind_for = 0
|
||||
|
||||
bond_unit_scale = self.config.get('bifrost_personality_bond_factor', 20000)
|
||||
self.num_peers = len(peers)
|
||||
num_peers = self.num_peers + 1e-10
|
||||
|
||||
self.tot_bond_factor = sum(
|
||||
p.get('encounters', 0) if isinstance(p, dict) else getattr(p, 'encounters', 0)
|
||||
for p in peers
|
||||
) / bond_unit_scale
|
||||
self.avg_bond_factor = self.tot_bond_factor / num_peers
|
||||
|
||||
num_aps_f = len(aps) + 1e-10
|
||||
num_sta = sum(len(ap.get('clients', [])) for ap in aps) + 1e-10
|
||||
aps_per_chan = [0.0] * NUM_CHANNELS
|
||||
sta_per_chan = [0.0] * NUM_CHANNELS
|
||||
peers_per_chan = [0.0] * NUM_CHANNELS
|
||||
|
||||
for ap in aps:
|
||||
ch_idx = ap.get('channel', 1) - 1
|
||||
if 0 <= ch_idx < NUM_CHANNELS:
|
||||
aps_per_chan[ch_idx] += 1.0
|
||||
sta_per_chan[ch_idx] += len(ap.get('clients', []))
|
||||
|
||||
for peer in peers:
|
||||
ch = peer.get('last_channel', 0) if isinstance(peer, dict) else getattr(peer, 'last_channel', 0)
|
||||
ch_idx = ch - 1
|
||||
if 0 <= ch_idx < NUM_CHANNELS:
|
||||
peers_per_chan[ch_idx] += 1.0
|
||||
|
||||
# Normalize
|
||||
aps_per_chan = [e / num_aps_f for e in aps_per_chan]
|
||||
sta_per_chan = [e / num_sta for e in sta_per_chan]
|
||||
peers_per_chan = [e / num_peers for e in peers_per_chan]
|
||||
|
||||
self._observation = {
|
||||
'aps_histogram': aps_per_chan,
|
||||
'sta_histogram': sta_per_chan,
|
||||
'peers_histogram': peers_per_chan,
|
||||
}
|
||||
self._observation_ready.set()
|
||||
|
||||
def track(self, deauth=False, assoc=False, handshake=False,
|
||||
hop=False, sleep=False, miss=False, inc=1):
|
||||
"""Increment epoch counters."""
|
||||
if deauth:
|
||||
self.num_deauths += inc
|
||||
self.did_deauth = True
|
||||
self.any_activity = True
|
||||
|
||||
if assoc:
|
||||
self.num_assocs += inc
|
||||
self.did_associate = True
|
||||
self.any_activity = True
|
||||
|
||||
if miss:
|
||||
self.num_missed += inc
|
||||
|
||||
if hop:
|
||||
self.num_hops += inc
|
||||
# Reset per-channel flags on hop
|
||||
self.did_deauth = False
|
||||
self.did_associate = False
|
||||
|
||||
if handshake:
|
||||
self.num_shakes += inc
|
||||
self.did_handshakes = True
|
||||
|
||||
if sleep:
|
||||
self.num_slept += inc
|
||||
|
||||
def next(self):
|
||||
"""Transition to next epoch — compute reward, update streaks, reset counters."""
|
||||
# Update activity streaks
|
||||
if not self.any_activity and not self.did_handshakes:
|
||||
self.inactive_for += 1
|
||||
self.active_for = 0
|
||||
else:
|
||||
self.active_for += 1
|
||||
self.inactive_for = 0
|
||||
self.sad_for = 0
|
||||
self.bored_for = 0
|
||||
|
||||
sad_threshold = self.config.get('bifrost_personality_sad_epochs', 25)
|
||||
bored_threshold = self.config.get('bifrost_personality_bored_epochs', 15)
|
||||
|
||||
if self.inactive_for >= sad_threshold:
|
||||
self.bored_for = 0
|
||||
self.sad_for += 1
|
||||
elif self.inactive_for >= bored_threshold:
|
||||
self.sad_for = 0
|
||||
self.bored_for += 1
|
||||
else:
|
||||
self.sad_for = 0
|
||||
self.bored_for = 0
|
||||
|
||||
now = time.time()
|
||||
self.epoch_duration = now - self.epoch_started
|
||||
|
||||
# System metrics
|
||||
cpu = _cpu_load()
|
||||
mem = _mem_usage()
|
||||
temp = _temperature()
|
||||
|
||||
# Cache epoch data for other threads
|
||||
self._epoch_data = {
|
||||
'duration_secs': self.epoch_duration,
|
||||
'slept_for_secs': self.num_slept,
|
||||
'blind_for_epochs': self.blind_for,
|
||||
'inactive_for_epochs': self.inactive_for,
|
||||
'active_for_epochs': self.active_for,
|
||||
'sad_for_epochs': self.sad_for,
|
||||
'bored_for_epochs': self.bored_for,
|
||||
'missed_interactions': self.num_missed,
|
||||
'num_hops': self.num_hops,
|
||||
'num_peers': self.num_peers,
|
||||
'tot_bond': self.tot_bond_factor,
|
||||
'avg_bond': self.avg_bond_factor,
|
||||
'num_deauths': self.num_deauths,
|
||||
'num_associations': self.num_assocs,
|
||||
'num_handshakes': self.num_shakes,
|
||||
'cpu_load': cpu,
|
||||
'mem_usage': mem,
|
||||
'temperature': temp,
|
||||
}
|
||||
self._epoch_data['reward'] = self._reward(self.epoch + 1, self._epoch_data)
|
||||
self._epoch_data_ready.set()
|
||||
|
||||
logger.info(
|
||||
"[epoch %d] dur=%ds blind=%d sad=%d bored=%d inactive=%d active=%d "
|
||||
"hops=%d missed=%d deauths=%d assocs=%d shakes=%d reward=%.3f",
|
||||
self.epoch, int(self.epoch_duration), self.blind_for,
|
||||
self.sad_for, self.bored_for, self.inactive_for, self.active_for,
|
||||
self.num_hops, self.num_missed, self.num_deauths, self.num_assocs,
|
||||
self.num_shakes, self._epoch_data['reward'],
|
||||
)
|
||||
|
||||
# Reset for next epoch
|
||||
self.epoch += 1
|
||||
self.epoch_started = now
|
||||
self.did_deauth = False
|
||||
self.num_deauths = 0
|
||||
self.num_peers = 0
|
||||
self.tot_bond_factor = 0.0
|
||||
self.avg_bond_factor = 0.0
|
||||
self.did_associate = False
|
||||
self.num_assocs = 0
|
||||
self.num_missed = 0
|
||||
self.did_handshakes = False
|
||||
self.num_shakes = 0
|
||||
self.num_hops = 0
|
||||
self.num_slept = 0
|
||||
self.any_activity = False
|
||||
|
||||
|
||||
# ── System metric helpers ────────────────────────────────────────
|
||||
|
||||
def _cpu_load():
|
||||
try:
|
||||
return os.getloadavg()[0]
|
||||
except (OSError, AttributeError):
|
||||
return 0.0
|
||||
|
||||
|
||||
def _mem_usage():
|
||||
try:
|
||||
with open('/proc/meminfo', 'r') as f:
|
||||
lines = f.readlines()
|
||||
total = int(lines[0].split()[1])
|
||||
available = int(lines[2].split()[1])
|
||||
return (total - available) / total if total else 0.0
|
||||
except Exception:
|
||||
return 0.0
|
||||
|
||||
|
||||
def _temperature():
|
||||
try:
|
||||
with open('/sys/class/thermal/thermal_zone0/temp', 'r') as f:
|
||||
return int(f.read().strip()) / 1000.0
|
||||
except Exception:
|
||||
return 0.0
|
||||
66
bifrost/faces.py
Normal file
66
bifrost/faces.py
Normal file
@@ -0,0 +1,66 @@
|
||||
"""
|
||||
Bifrost — ASCII face definitions.
|
||||
Ported from pwnagotchi/ui/faces.py with full face set.
|
||||
"""
|
||||
|
||||
LOOK_R = '( \u2686_\u2686)'
|
||||
LOOK_L = '(\u2609_\u2609 )'
|
||||
LOOK_R_HAPPY = '( \u25d5\u203f\u25d5)'
|
||||
LOOK_L_HAPPY = '(\u25d5\u203f\u25d5 )'
|
||||
SLEEP = '(\u21c0\u203f\u203f\u21bc)'
|
||||
SLEEP2 = '(\u2256\u203f\u203f\u2256)'
|
||||
AWAKE = '(\u25d5\u203f\u203f\u25d5)'
|
||||
BORED = '(-__-)'
|
||||
INTENSE = '(\u00b0\u25c3\u25c3\u00b0)'
|
||||
COOL = '(\u2310\u25a0_\u25a0)'
|
||||
HAPPY = '(\u2022\u203f\u203f\u2022)'
|
||||
GRATEFUL = '(^\u203f\u203f^)'
|
||||
EXCITED = '(\u1d54\u25e1\u25e1\u1d54)'
|
||||
MOTIVATED = '(\u263c\u203f\u203f\u263c)'
|
||||
DEMOTIVATED = '(\u2256__\u2256)'
|
||||
SMART = '(\u271c\u203f\u203f\u271c)'
|
||||
LONELY = '(\u0628__\u0628)'
|
||||
SAD = '(\u2565\u2601\u2565 )'
|
||||
ANGRY = "(-_-')"
|
||||
FRIEND = '(\u2665\u203f\u203f\u2665)'
|
||||
BROKEN = '(\u2613\u203f\u203f\u2613)'
|
||||
DEBUG = '(#__#)'
|
||||
UPLOAD = '(1__0)'
|
||||
UPLOAD1 = '(1__1)'
|
||||
UPLOAD2 = '(0__1)'
|
||||
STARTING = '(. .)'
|
||||
READY = '( ^_^)'
|
||||
|
||||
# Map mood name → face constant
|
||||
MOOD_FACES = {
|
||||
'starting': STARTING,
|
||||
'ready': READY,
|
||||
'sleeping': SLEEP,
|
||||
'awake': AWAKE,
|
||||
'bored': BORED,
|
||||
'sad': SAD,
|
||||
'angry': ANGRY,
|
||||
'excited': EXCITED,
|
||||
'lonely': LONELY,
|
||||
'grateful': GRATEFUL,
|
||||
'happy': HAPPY,
|
||||
'cool': COOL,
|
||||
'intense': INTENSE,
|
||||
'motivated': MOTIVATED,
|
||||
'demotivated': DEMOTIVATED,
|
||||
'friend': FRIEND,
|
||||
'broken': BROKEN,
|
||||
'debug': DEBUG,
|
||||
'smart': SMART,
|
||||
}
|
||||
|
||||
|
||||
def load_from_config(config):
|
||||
"""Override faces from config dict (e.g. custom emojis)."""
|
||||
for face_name, face_value in (config or {}).items():
|
||||
key = face_name.upper()
|
||||
if key in globals():
|
||||
globals()[key] = face_value
|
||||
lower = face_name.lower()
|
||||
if lower in MOOD_FACES:
|
||||
MOOD_FACES[lower] = face_value
|
||||
198
bifrost/plugins.py
Normal file
198
bifrost/plugins.py
Normal file
@@ -0,0 +1,198 @@
|
||||
"""
|
||||
Bifrost — Plugin system.
|
||||
Ported from pwnagotchi/plugins/__init__.py with ThreadPoolExecutor.
|
||||
Compatible with existing pwnagotchi plugin files.
|
||||
"""
|
||||
import os
|
||||
import glob
|
||||
import threading
|
||||
import importlib
|
||||
import importlib.util
|
||||
import logging
|
||||
import concurrent.futures
|
||||
|
||||
from logger import Logger
|
||||
|
||||
logger = Logger(name="bifrost.plugins", level=logging.DEBUG)
|
||||
|
||||
default_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), "plugins")
|
||||
loaded = {}
|
||||
database = {}
|
||||
locks = {}
|
||||
|
||||
_executor = concurrent.futures.ThreadPoolExecutor(
|
||||
max_workers=4, thread_name_prefix="BifrostPlugin"
|
||||
)
|
||||
|
||||
|
||||
class Plugin:
|
||||
"""Base class for Bifrost/Pwnagotchi plugins.
|
||||
|
||||
Subclasses are auto-registered via __init_subclass__.
|
||||
"""
|
||||
__author__ = 'unknown'
|
||||
__version__ = '0.0.0'
|
||||
__license__ = 'GPL3'
|
||||
__description__ = ''
|
||||
__name__ = ''
|
||||
__help__ = ''
|
||||
__dependencies__ = []
|
||||
__defaults__ = {}
|
||||
|
||||
def __init__(self):
|
||||
self.options = {}
|
||||
|
||||
@classmethod
|
||||
def __init_subclass__(cls, **kwargs):
|
||||
super().__init_subclass__(**kwargs)
|
||||
global loaded, locks
|
||||
|
||||
plugin_name = cls.__module__.split('.')[0]
|
||||
plugin_instance = cls()
|
||||
logger.debug("loaded plugin %s as %s", plugin_name, plugin_instance)
|
||||
loaded[plugin_name] = plugin_instance
|
||||
|
||||
for attr_name in dir(plugin_instance):
|
||||
if attr_name.startswith('on_'):
|
||||
cb = getattr(plugin_instance, attr_name, None)
|
||||
if cb is not None and callable(cb):
|
||||
locks["%s::%s" % (plugin_name, attr_name)] = threading.Lock()
|
||||
|
||||
|
||||
def toggle_plugin(name, enable=True):
|
||||
"""Enable or disable a plugin at runtime. Returns True if state changed."""
|
||||
global loaded, database
|
||||
|
||||
if not enable and name in loaded:
|
||||
try:
|
||||
if hasattr(loaded[name], 'on_unload'):
|
||||
loaded[name].on_unload()
|
||||
except Exception as e:
|
||||
logger.warning("Error unloading plugin %s: %s", name, e)
|
||||
del loaded[name]
|
||||
return True
|
||||
|
||||
if enable and name in database and name not in loaded:
|
||||
try:
|
||||
load_from_file(database[name])
|
||||
if name in loaded:
|
||||
one(name, 'loaded')
|
||||
return True
|
||||
except Exception as e:
|
||||
logger.warning("Error loading plugin %s: %s", name, e)
|
||||
|
||||
return False
|
||||
|
||||
|
||||
def on(event_name, *args, **kwargs):
|
||||
"""Dispatch event to ALL loaded plugins."""
|
||||
for plugin_name in list(loaded.keys()):
|
||||
one(plugin_name, event_name, *args, **kwargs)
|
||||
|
||||
|
||||
def _locked_cb(lock_name, cb, *args, **kwargs):
|
||||
"""Execute callback under its per-plugin lock."""
|
||||
global locks
|
||||
if lock_name not in locks:
|
||||
locks[lock_name] = threading.Lock()
|
||||
with locks[lock_name]:
|
||||
cb(*args, **kwargs)
|
||||
|
||||
|
||||
def one(plugin_name, event_name, *args, **kwargs):
|
||||
"""Dispatch event to a single plugin (thread-safe)."""
|
||||
global loaded
|
||||
if plugin_name in loaded:
|
||||
plugin = loaded[plugin_name]
|
||||
cb_name = 'on_%s' % event_name
|
||||
callback = getattr(plugin, cb_name, None)
|
||||
if callback is not None and callable(callback):
|
||||
try:
|
||||
lock_name = "%s::%s" % (plugin_name, cb_name)
|
||||
_executor.submit(_locked_cb, lock_name, callback, *args, **kwargs)
|
||||
except Exception as e:
|
||||
logger.error("error running %s.%s: %s", plugin_name, cb_name, e)
|
||||
|
||||
|
||||
def load_from_file(filename):
|
||||
"""Load a single plugin file."""
|
||||
logger.debug("loading %s", filename)
|
||||
plugin_name = os.path.basename(filename.replace(".py", ""))
|
||||
spec = importlib.util.spec_from_file_location(plugin_name, filename)
|
||||
instance = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(instance)
|
||||
return plugin_name, instance
|
||||
|
||||
|
||||
def load_from_path(path, enabled=()):
|
||||
"""Scan a directory for plugins, load enabled ones."""
|
||||
global loaded, database
|
||||
if not path or not os.path.isdir(path):
|
||||
return loaded
|
||||
|
||||
logger.debug("loading plugins from %s — enabled: %s", path, enabled)
|
||||
for filename in glob.glob(os.path.join(path, "*.py")):
|
||||
plugin_name = os.path.basename(filename.replace(".py", ""))
|
||||
database[plugin_name] = filename
|
||||
if plugin_name in enabled:
|
||||
try:
|
||||
load_from_file(filename)
|
||||
except Exception as e:
|
||||
logger.warning("error loading %s: %s", filename, e)
|
||||
|
||||
return loaded
|
||||
|
||||
|
||||
def load(config):
|
||||
"""Load plugins from default + custom paths based on config."""
|
||||
plugins_cfg = config.get('bifrost_plugins', {})
|
||||
enabled = [
|
||||
name for name, opts in plugins_cfg.items()
|
||||
if isinstance(opts, dict) and opts.get('enabled', False)
|
||||
]
|
||||
|
||||
# Load from default path (bifrost/plugins/)
|
||||
if os.path.isdir(default_path):
|
||||
load_from_path(default_path, enabled=enabled)
|
||||
|
||||
# Load from custom path
|
||||
custom_path = config.get('bifrost_plugins_path', '')
|
||||
if custom_path and os.path.isdir(custom_path):
|
||||
load_from_path(custom_path, enabled=enabled)
|
||||
|
||||
# Propagate options
|
||||
for name, plugin in loaded.items():
|
||||
if name in plugins_cfg:
|
||||
plugin.options = plugins_cfg[name]
|
||||
|
||||
on('loaded')
|
||||
on('config_changed', config)
|
||||
|
||||
|
||||
def get_loaded_info():
|
||||
"""Return list of loaded plugin info dicts for web API."""
|
||||
result = []
|
||||
for name, plugin in loaded.items():
|
||||
result.append({
|
||||
'name': name,
|
||||
'enabled': True,
|
||||
'author': getattr(plugin, '__author__', 'unknown'),
|
||||
'version': getattr(plugin, '__version__', '0.0.0'),
|
||||
'description': getattr(plugin, '__description__', ''),
|
||||
})
|
||||
# Also include known-but-not-loaded plugins
|
||||
for name, path in database.items():
|
||||
if name not in loaded:
|
||||
result.append({
|
||||
'name': name,
|
||||
'enabled': False,
|
||||
'author': '',
|
||||
'version': '',
|
||||
'description': '',
|
||||
})
|
||||
return result
|
||||
|
||||
|
||||
def shutdown():
|
||||
"""Clean shutdown of plugin system."""
|
||||
_executor.shutdown(wait=False)
|
||||
155
bifrost/voice.py
Normal file
155
bifrost/voice.py
Normal file
@@ -0,0 +1,155 @@
|
||||
"""
|
||||
Bifrost — Voice / status messages.
|
||||
Ported from pwnagotchi/voice.py, uses random choice for personality.
|
||||
"""
|
||||
import random
|
||||
|
||||
|
||||
class BifrostVoice:
|
||||
"""Returns random contextual messages for the Bifrost UI."""
|
||||
|
||||
def on_starting(self):
|
||||
return random.choice([
|
||||
"Hi, I'm Bifrost! Starting ...",
|
||||
"New day, new hunt, new pwns!",
|
||||
"Hack the Planet!",
|
||||
"Initializing WiFi recon ...",
|
||||
])
|
||||
|
||||
def on_ready(self):
|
||||
return random.choice([
|
||||
"Ready to roll!",
|
||||
"Let's find some handshakes!",
|
||||
"WiFi recon active.",
|
||||
])
|
||||
|
||||
def on_ai_ready(self):
|
||||
return random.choice([
|
||||
"AI ready.",
|
||||
"The neural network is ready.",
|
||||
])
|
||||
|
||||
def on_normal(self):
|
||||
return random.choice(['', '...'])
|
||||
|
||||
def on_free_channel(self, channel):
|
||||
return f"Hey, channel {channel} is free!"
|
||||
|
||||
def on_bored(self):
|
||||
return random.choice([
|
||||
"I'm bored ...",
|
||||
"Let's go for a walk!",
|
||||
"Nothing interesting around here ...",
|
||||
])
|
||||
|
||||
def on_motivated(self, reward):
|
||||
return "This is the best day of my life!"
|
||||
|
||||
def on_demotivated(self, reward):
|
||||
return "Shitty day :/"
|
||||
|
||||
def on_sad(self):
|
||||
return random.choice([
|
||||
"I'm extremely bored ...",
|
||||
"I'm very sad ...",
|
||||
"I'm sad",
|
||||
"...",
|
||||
])
|
||||
|
||||
def on_angry(self):
|
||||
return random.choice([
|
||||
"...",
|
||||
"Leave me alone ...",
|
||||
"I'm mad at you!",
|
||||
])
|
||||
|
||||
def on_excited(self):
|
||||
return random.choice([
|
||||
"I'm living the life!",
|
||||
"I pwn therefore I am.",
|
||||
"So many networks!!!",
|
||||
"I'm having so much fun!",
|
||||
"My crime is that of curiosity ...",
|
||||
])
|
||||
|
||||
def on_new_peer(self, peer_name, first_encounter=False):
|
||||
if first_encounter:
|
||||
return f"Hello {peer_name}! Nice to meet you."
|
||||
return random.choice([
|
||||
f"Yo {peer_name}! Sup?",
|
||||
f"Hey {peer_name} how are you doing?",
|
||||
f"Unit {peer_name} is nearby!",
|
||||
])
|
||||
|
||||
def on_lost_peer(self, peer_name):
|
||||
return random.choice([
|
||||
f"Uhm ... goodbye {peer_name}",
|
||||
f"{peer_name} is gone ...",
|
||||
])
|
||||
|
||||
def on_miss(self, who):
|
||||
return random.choice([
|
||||
f"Whoops ... {who} is gone.",
|
||||
f"{who} missed!",
|
||||
"Missed!",
|
||||
])
|
||||
|
||||
def on_grateful(self):
|
||||
return random.choice([
|
||||
"Good friends are a blessing!",
|
||||
"I love my friends!",
|
||||
])
|
||||
|
||||
def on_lonely(self):
|
||||
return random.choice([
|
||||
"Nobody wants to play with me ...",
|
||||
"I feel so alone ...",
|
||||
"Where's everybody?!",
|
||||
])
|
||||
|
||||
def on_napping(self, secs):
|
||||
return random.choice([
|
||||
f"Napping for {secs}s ...",
|
||||
"Zzzzz",
|
||||
f"ZzzZzzz ({secs}s)",
|
||||
])
|
||||
|
||||
def on_shutdown(self):
|
||||
return random.choice(["Good night.", "Zzz"])
|
||||
|
||||
def on_awakening(self):
|
||||
return random.choice(["...", "!"])
|
||||
|
||||
def on_waiting(self, secs):
|
||||
return random.choice([
|
||||
f"Waiting for {secs}s ...",
|
||||
"...",
|
||||
f"Looking around ({secs}s)",
|
||||
])
|
||||
|
||||
def on_assoc(self, ap_name):
|
||||
return random.choice([
|
||||
f"Hey {ap_name} let's be friends!",
|
||||
f"Associating to {ap_name}",
|
||||
f"Yo {ap_name}!",
|
||||
])
|
||||
|
||||
def on_deauth(self, sta_mac):
|
||||
return random.choice([
|
||||
f"Just decided that {sta_mac} needs no WiFi!",
|
||||
f"Deauthenticating {sta_mac}",
|
||||
f"Kickbanning {sta_mac}!",
|
||||
])
|
||||
|
||||
def on_handshakes(self, new_shakes):
|
||||
s = 's' if new_shakes > 1 else ''
|
||||
return f"Cool, we got {new_shakes} new handshake{s}!"
|
||||
|
||||
def on_rebooting(self):
|
||||
return "Oops, something went wrong ... Rebooting ..."
|
||||
|
||||
def on_epoch(self, epoch_num):
|
||||
return random.choice([
|
||||
f"Epoch {epoch_num} complete.",
|
||||
f"Finished epoch {epoch_num}.",
|
||||
])
|
||||
Reference in New Issue
Block a user