mirror of
https://github.com/infinition/Bjorn.git
synced 2026-03-16 01: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:
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
|
||||
Reference in New Issue
Block a user