import argparse import json import os import threading import time from typing import List import redis import requests from flask import jsonify, request, Flask from flask_caching import Cache from urllib3.exceptions import InsecureRequestWarning from checker.units import filesize requests.packages.urllib3.disable_warnings(category=InsecureRequestWarning) MAX_POINTS_PER_IP = 1900 OPNSENSE_URL = os.environ.get('OPN_URL') OPNSENSE_KEY = os.environ.get('OPN_KEY') OPNSENSE_SECRET = os.environ.get('OPN_SECRET') if not OPNSENSE_URL or not OPNSENSE_KEY or not OPNSENSE_URL: raise Exception('Missing environment variables') OPNSENSE_URL = OPNSENSE_URL.strip('/') class TrafficEntry: def __init__(self, interface: str, address: str, rate_bits_in: int, rate_bits_out: int, rate_bits: int, cumulative_bytes_in: int, cumulative_bytes_out: int, cumulative_bytes: int, connections: dict, timestamp: float): self.interface = interface self.address = address self.rate_bits_in = rate_bits_in self.rate_bits_out = rate_bits_out self.rate_bits = rate_bits self.cumulative_bytes_in = cumulative_bytes_in self.cumulative_bytes_out = cumulative_bytes_out self.cumulative_bytes = cumulative_bytes self.connections = connections self.timestamp = timestamp def to_json(self): return self.__dict__ @classmethod def from_json(cls, json_str): data = json.loads(json_str) return cls(**data) class OpnsenseTraffic: def __init__(self): self.redis = redis.Redis(host='localhost', port=6379, db=1) def flush(self): self.redis.flushdb() def add_entry(self, item: TrafficEntry): # TODO: kick out the oldest item key = f"{item.interface}:{item.address}" if self.redis.llen(key) >= MAX_POINTS_PER_IP: self.redis.lpop(key) self.redis.rpush(key, json.dumps(item.to_json())) def get_address(self, input_address: str) -> List[TrafficEntry]: keys = self.redis.keys(f"*:{input_address}") data = [] for key in keys: entries = self.redis.lrange(key, 0, -1) data.extend([TrafficEntry.from_json(entry.decode()) for entry in entries]) return data def get_entries(self, input_address: str): keys = self.redis.keys() data = {} for key in keys: try: interface, address = key.decode().split(":") except ValueError: # Can get things like "opt2:::" continue if address != input_address: continue entries = self.redis.lrange(key, 0, -1) if interface not in data: data[interface] = {} data[interface][address] = [TrafficEntry.from_json(entry.decode()).to_json() for entry in entries] return data def get_traffic(self, address: str, minus_seconds: int = 0, human: bool = False): max_rate_in = 0 max_rate_out = 0 bytes_in = 0 bytes_out = 0 connections = 0 if minus_seconds == 0: minus_sec_diff = 0 else: minus_sec_diff = int(time.time()) - minus_seconds address_traffic = self.get_address(address) for entry in address_traffic: if entry.timestamp >= minus_sec_diff: max_rate_in = max(max_rate_in, entry.rate_bits_in) max_rate_out = max(max_rate_out, entry.rate_bits_out) bytes_in += entry.cumulative_bytes_in bytes_out += entry.cumulative_bytes_out connections += len(entry.connections) if human: return filesize(max_rate_in), filesize(max_rate_out), filesize(bytes_in), filesize(bytes_out), connections else: return max_rate_in, max_rate_out, bytes_in, bytes_out, connections def get_interfaces(): r = redis.Redis(host='localhost', port=6379, db=2) try: return json.loads(r.get('interfaces')) except Exception as e: return [] def get_interface_names(): r = redis.Redis(host='localhost', port=6379, db=2) # Map interface names to their internal names while True: interfaces_mapping_response = requests.get(f'{OPNSENSE_URL}/api/diagnostics/traffic/interface', headers={'Accept': 'application/json'}, auth=(OPNSENSE_KEY, OPNSENSE_SECRET), verify=False) interfaces_mapping_response.raise_for_status() interfaces = list(interfaces_mapping_response.json()['interfaces'].keys()) if 'interface' in interfaces: interfaces.remove('interface') # Store the interfaces in Redis r.set('interfaces', json.dumps(interfaces)) time.sleep(60) def background_thread(): traffic_data = OpnsenseTraffic() traffic_data.flush() while True: start_time = time.time() interface_req = ','.join(get_interfaces()) response = requests.get(f'{OPNSENSE_URL}/api/diagnostics/traffic/top/{interface_req}', headers={'Accept': 'application/json'}, auth=(OPNSENSE_KEY, OPNSENSE_SECRET), verify=False) response.raise_for_status() timestamp = time.time() for interface, data in response.json().items(): for item in data.get('records'): traffic_data.add_entry( TrafficEntry( address=item['address'], interface=interface, rate_bits=item['rate_bits'], rate_bits_in=item['rate_bits_in'], rate_bits_out=item['rate_bits_out'], cumulative_bytes=item['cumulative_bytes'], cumulative_bytes_in=item['cumulative_bytes_in'], cumulative_bytes_out=item['cumulative_bytes_out'], connections=item['details'], timestamp=timestamp ) ) end_time = time.time() api_request_time = end_time - start_time adjusted_sleep_duration = max(1 - api_request_time, 0) time.sleep(adjusted_sleep_duration) flask_traffic_data = OpnsenseTraffic() app = Flask(__name__) cache = Cache(app, config={ "CACHE_TYPE": "RedisCache", "CACHE_REDIS_HOST": "127.0.0.1", "port": 6379 }) @app.route('/traffic/
', methods=['GET']) @app.route('/traffic/
/', methods=['GET']) @cache.cached(timeout=10, query_string=True) def get_traffic(address, interface=None): minus_seconds = request.args.get('seconds', default=0, type=int) human = request.args.get('human', default=False, type=bool) max_rate_in, max_rate_out, bytes_in, bytes_out, connections = flask_traffic_data.get_traffic(address, minus_seconds, human) entries = flask_traffic_data.get_entries(address) num_entries = 0 if interface: if interface not in entries.keys(): return 'Interface not found for address', 404 entries = entries[interface] num_entries = len(entries) return jsonify({ interface: { 'max_rate_in': max_rate_in, 'max_rate_out': max_rate_out, 'bytes_in': bytes_in, 'bytes_out': bytes_out, 'connections': connections }, 'entries': num_entries }) else: for interface in entries: num_entries += len(entries[interface][list(entries[interface].keys())[0]]) return jsonify({ 'data': { 'max_rate_in': max_rate_in, 'max_rate_out': max_rate_out, 'bytes_in': bytes_in, 'bytes_out': bytes_out, 'connections': connections }, 'entries': num_entries }) @app.route('/data/', methods=['GET']) @cache.cached(timeout=10) def get_entries(ip): entries = flask_traffic_data.get_entries(ip) # Remove the IP address from the dict new_entries = {} for key, value in entries.items(): new_entries[key] = [] for sub_key, sub_value in value.items(): new_entries[key].extend(sub_value) return jsonify(new_entries) @app.route('/interfaces', methods=['GET']) @cache.cached(timeout=10) def flask_get_interfaces(): return jsonify(get_interfaces()) if __name__ == '__main__': parser = argparse.ArgumentParser() parser.add_argument('--daemon', action='store_true', help='Run the background daemon.') args = parser.parse_args() if args.daemon: t1 = threading.Thread(target=get_interface_names, daemon=True) t1.start() print('Fetching interface list... ', end='') while not len(get_interfaces()): time.sleep(2) print('Done!') t2 = threading.Thread(target=background_thread, daemon=True) t2.start() print('Threads started!') while True: time.sleep(10000) else: app.run(host='0.0.0.0', debug=True)