PacketSentry_Linux/main.py at main · PacketSentry/PacketSentry_Linux · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
import psutil
from kivy.config import Config
Config.set('input', 'mouse', 'mouse,multitouch_on_demand')

from kivy.app import App
from kivy.lang import Builder
from kivy.clock import Clock
from kivy.core.window import Window
from kivy.utils import platform

# Import the UI (existing files)
from ui.widgets import TrafficGraph, AppDashboard, PingGraph, LoginPopup

# IMPORT RUST CORE
# Note: You must build this first using 'maturin develop'
import packet_sentry_core

class NetworkApp(App):
    def build(self):
        Window.size = (900, 700)
        return Builder.load_file("ui/dashboard.kv")

    def on_start(self):
        # 1. Initialize Rust Backend
        self.backend = packet_sentry_core.CoreBackend()
        self.backend.start()
        print("Rust Core Started...")

        # 2. UI State
        self.accumulated_rates = {} # To smooth out graph

        # 3. Schedule Updates
        Clock.schedule_interval(self.update_ui, 1.0)

        # Hardware counters for total graph verification
        self.last_net_io = psutil.net_io_counters()

    def update_ui(self, dt):
        # --- A. Get Data from Rust ---
        # Returns dict: {'Firefox': (down_bytes, up_bytes), ...}
        traffic_delta = self.backend.get_traffic_update()
        pings = self.backend.get_pings()

        # --- B. Calculate Rates (Bytes per second) ---
        # Since Rust clears the buffer on read, 'traffic_delta' IS the rate per second (approx)

        # Prepare data for AppDashboard
        # Dashboard expects: {'App': [down_spd_kb, up_spd_kb]}
        ui_rates = {}
        total_down = 0.0
        total_up = 0.0

        for app, (bytes_down, bytes_up) in traffic_delta.items():
            d_kb = bytes_down / 1024.0
            u_kb = bytes_up / 1024.0

            ui_rates[app] = [d_kb, u_kb]
            total_down += d_kb
            total_up += u_kb

        # --- C. Update UI Widgets ---

        # 1. Main Graph (Hardware Stats fallback for total accuracy)
        current_net_io = psutil.net_io_counters()
        hw_down = (current_net_io.bytes_recv - self.last_net_io.bytes_recv) / 1024
        hw_up = (current_net_io.bytes_sent - self.last_net_io.bytes_sent) / 1024
        self.last_net_io = current_net_io

        if "main_graph" in self.root.ids:
            # Use Hardware stats for the main graph as they are absolute
            self.root.ids.main_graph.update_graph(hw_down, hw_up)

        # 2. App Dashboard (Uses Rust Data)
        if "dashboard" in self.root.ids:
            self.root.ids.dashboard.update_apps(ui_rates)

        # 3. Ping Graph (Uses Rust Data)
        if "ping_graph" in self.root.ids:
            cf = pings.get("Cloudflare", 0.0)
            gg = pings.get("Google", 0.0)
            self.root.ids.ping_graph.update_graph(cf, gg)

    def on_stop(self):
        if hasattr(self, 'backend'):
            self.backend.stop()

if __name__ == "__main__":
    NetworkApp().run()