mirror of
https://github.com/infinition/Bjorn.git
synced 2026-03-15 08:52:00 +00:00
aac77a3e76
- 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.
293 lines
9.4 KiB
Python
293 lines
9.4 KiB
Python
"""
|
|
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
|