varun-r-mallya/python-bpf
1
0
Fork 0
mirror of https://github.com/varun-r-mallya/Python-BPF.git synced 2025-12-31 21:06:25 +00:00
Code Issues Packages Projects Releases Wiki Activity

format chore

Browse Source
This commit is contained in:
varun-r-mallya
2025-11-28 20:59:46 +05:30
parent bbe4990878
commit d73c793989
3 changed files with 98 additions and 64 deletions
Download Patch File Download Diff File Expand all files Collapse all files

27
BCC-Examples/container-monitor/container_monitor.py
View File

@ -1,20 +1,18 @@
"""Container Monitor - TUI-based cgroup monitoring combining syscall, file I/O, and network tracking.""" """Container Monitor - TUI-based cgroup monitoring combining syscall, file I/O, and network tracking."""
import time
import os
from pathlib import Path
from pythonbpf import bpf, map, section, bpfglobal, struct, BPF from pythonbpf import bpf, map, section, bpfglobal, struct, BPF
from pythonbpf.maps import HashMap from pythonbpf.maps import HashMap
from pythonbpf.helper import get_current_cgroup_id from pythonbpf.helper import get_current_cgroup_id
from ctypes import c_int32, c_uint64, c_void_p from ctypes import c_int32, c_uint64, c_void_p
from vmlinux import struct_pt_regs, struct_sk_buff from vmlinux import struct_pt_regs, struct_sk_buff
from data_collector import ContainerDataCollector from data_collection import ContainerDataCollector
from tui import ContainerMonitorTUI from tui import ContainerMonitorTUI
# ==================== BPF Structs ==================== # ==================== BPF Structs ====================
@bpf @bpf
@struct @struct
class read_stats: class read_stats:
@ -40,6 +38,7 @@ class net_stats:
# ==================== BPF Maps ==================== # ==================== BPF Maps ====================
@bpf @bpf
@map @map
def read_map() -> HashMap: def read_map() -> HashMap:
@ -66,6 +65,7 @@ def syscall_count() -> HashMap:
# ==================== File I/O Tracing ==================== # ==================== File I/O Tracing ====================
@bpf @bpf
@section("kprobe/vfs_read") @section("kprobe/vfs_read")
def trace_read(ctx: struct_pt_regs) -> c_int32: def trace_read(ctx: struct_pt_regs) -> c_int32:
@ -109,6 +109,7 @@ def trace_write(ctx1: struct_pt_regs) -> c_int32:
# ==================== Network I/O Tracing ==================== # ==================== Network I/O Tracing ====================
@bpf @bpf
@section("kprobe/__netif_receive_skb") @section("kprobe/__netif_receive_skb")
def trace_netif_rx(ctx2: struct_pt_regs) -> c_int32: def trace_netif_rx(ctx2: struct_pt_regs) -> c_int32:
@ -165,6 +166,7 @@ def trace_dev_xmit(ctx3: struct_pt_regs) -> c_int32:
# ==================== Syscall Tracing ==================== # ==================== Syscall Tracing ====================
@bpf @bpf
@section("tracepoint/raw_syscalls/sys_enter") @section("tracepoint/raw_syscalls/sys_enter")
def count_syscalls(ctx: c_void_p) -> c_int32: def count_syscalls(ctx: c_void_p) -> c_int32:
@ -190,32 +192,29 @@ def LICENSE() -> str:
if __name__ == "__main__": if __name__ == "__main__":
print("🔥 Loading BPF programs...") print("🔥 Loading BPF programs...")
# Load and attach BPF program # Load and attach BPF program
b = BPF() b = BPF()
b.load() b.load()
b.attach_all() b.attach_all()
# Get map references and enable struct deserialization # Get map references and enable struct deserialization
read_map_ref = b["read_map"] read_map_ref = b["read_map"]
write_map_ref = b["write_map"] write_map_ref = b["write_map"]
net_stats_map_ref = b["net_stats_map"] net_stats_map_ref = b["net_stats_map"]
syscall_count_ref = b["syscall_count"] syscall_count_ref = b["syscall_count"]
read_map_ref.set_value_struct("read_stats") read_map_ref.set_value_struct("read_stats")
write_map_ref.set_value_struct("write_stats") write_map_ref.set_value_struct("write_stats")
net_stats_map_ref.set_value_struct("net_stats") net_stats_map_ref.set_value_struct("net_stats")
print("✅ BPF programs loaded and attached") print("✅ BPF programs loaded and attached")
# Setup data collector # Setup data collector
collector = ContainerDataCollector( collector = ContainerDataCollector(
read_map_ref, read_map_ref, write_map_ref, net_stats_map_ref, syscall_count_ref
write_map_ref,
net_stats_map_ref,
syscall_count_ref
) )
# Create and run TUI # Create and run TUI
tui = ContainerMonitorTUI(collector) tui = ContainerMonitorTUI(collector)
tui.run() tui.run()

28
BCC-Examples/container-monitor/data_collector.py → BCC-Examples/container-monitor/data_collection.py
View File

@ -3,7 +3,7 @@
import os import os
import time import time
from pathlib import Path from pathlib import Path
from typing import Dict, List, Set, Optional, Tuple from typing import Dict, List, Set, Optional
from dataclasses import dataclass from dataclasses import dataclass
from collections import deque, defaultdict from collections import deque, defaultdict
@ -11,6 +11,7 @@ from collections import deque, defaultdict
@dataclass @dataclass
class CgroupInfo: class CgroupInfo:
"""Information about a cgroup.""" """Information about a cgroup."""
id: int id: int
name: str name: str
path: str path: str
@ -19,6 +20,7 @@ class CgroupInfo:
@dataclass @dataclass
class ContainerStats: class ContainerStats:
"""Statistics for a container/cgroup.""" """Statistics for a container/cgroup."""
cgroup_id: int cgroup_id: int
cgroup_name: str cgroup_name: str
@ -44,7 +46,9 @@ class ContainerStats:
class ContainerDataCollector: class ContainerDataCollector:
"""Collects and manages container monitoring data from BPF.""" """Collects and manages container monitoring data from BPF."""
def __init__(self, read_map, write_map, net_stats_map, syscall_map, history_size: int = 100): def __init__(
self, read_map, write_map, net_stats_map, syscall_map, history_size: int = 100
):
self.read_map = read_map self.read_map = read_map
self.write_map = write_map self.write_map = write_map
self.net_stats_map = net_stats_map self.net_stats_map = net_stats_map
@ -53,12 +57,14 @@ class ContainerDataCollector:
# Caching # Caching
self._cgroup_cache: Dict[int, CgroupInfo] = {} self._cgroup_cache: Dict[int, CgroupInfo] = {}
self._cgroup_cache_time = 0 self._cgroup_cache_time = 0
self._cache_ttl = 5. self._cache_ttl = 5.0
0 # Refresh cache every 5 seconds 0 # Refresh cache every 5 seconds
# Historical data for graphing # Historical data for graphing
self._history_size = history_size self._history_size = history_size
self._history: Dict[int, deque] = defaultdict(lambda: deque(maxlen=history_size)) self._history: Dict[int, deque] = defaultdict(
lambda: deque(maxlen=history_size)
)
def get_all_cgroups(self) -> List[CgroupInfo]: def get_all_cgroups(self) -> List[CgroupInfo]:
"""Get all cgroups with caching.""" """Get all cgroups with caching."""
@ -105,11 +111,7 @@ class ContainerDataCollector:
best_path = self._get_best_cgroup_path(paths) best_path = self._get_best_cgroup_path(paths)
name = self._get_cgroup_name(best_path) name = self._get_cgroup_name(best_path)
new_cache[cgroup_id] = CgroupInfo( new_cache[cgroup_id] = CgroupInfo(id=cgroup_id, name=name, path=best_path)
id=cgroup_id,
name=name,
path=best_path
)
self._cgroup_cache = new_cache self._cgroup_cache = new_cache
self._cgroup_cache_time = time.time() self._cgroup_cache_time = time.time()
@ -120,13 +122,13 @@ class ContainerDataCollector:
# Prefer paths with more components (more specific) # Prefer paths with more components (more specific)
# Prefer paths containing docker, podman, etc. # Prefer paths containing docker, podman, etc.
for keyword in ['docker', 'podman', 'kubernetes', 'k8s', 'systemd']: for keyword in ["docker", "podman", "kubernetes", "k8s", "systemd"]:
for path in path_list: for path in path_list:
if keyword in path.lower(): if keyword in path.lower():
return path return path
# Return longest path (most specific) # Return longest path (most specific)
return max(path_list, key=lambda p: (len(p.split('/')), len(p))) return max(path_list, key=lambda p: (len(p.split("/")), len(p)))
def _get_cgroup_name(self, path: str) -> str: def _get_cgroup_name(self, path: str) -> str:
"""Extract a friendly name from cgroup path.""" """Extract a friendly name from cgroup path."""
@ -165,9 +167,7 @@ class ContainerDataCollector:
cgroup_name = cgroup_info.name if cgroup_info else f"cgroup-{cgroup_id}" cgroup_name = cgroup_info.name if cgroup_info else f"cgroup-{cgroup_id}"
stats = ContainerStats( stats = ContainerStats(
cgroup_id=cgroup_id, cgroup_id=cgroup_id, cgroup_name=cgroup_name, timestamp=time.time()
cgroup_name=cgroup_name,
timestamp=time.time()
) )
# Get file I/O stats # Get file I/O stats

107
BCC-Examples/container-monitor/tui.py
View File

@ -1,10 +1,9 @@
"""Terminal User Interface for container monitoring.""" """Terminal User Interface for container monitoring."""
import sys
import time import time
import curses import curses
from typing import Optional, List from typing import Optional, List
from data_collector import ContainerDataCollector, CgroupInfo, ContainerStats from data_collection import ContainerDataCollector
class ContainerMonitorTUI: class ContainerMonitorTUI:
@ -122,11 +121,11 @@ class ContainerMonitorTUI:
# Highlight selected # Highlight selected
stdscr.attron(curses.color_pair(2) | curses.A_BOLD | curses.A_REVERSE) stdscr.attron(curses.color_pair(2) | curses.A_BOLD | curses.A_REVERSE)
line = f"► {cgroup.name:<40} ID: {cgroup.id}" line = f"► {cgroup.name:<40} ID: {cgroup.id}"
stdscr.addstr(y, 2, line[:width - 4]) stdscr.addstr(y, 2, line[: width - 4])
stdscr.attroff(curses.color_pair(2) | curses.A_BOLD | curses.A_REVERSE) stdscr.attroff(curses.color_pair(2) | curses.A_BOLD | curses.A_REVERSE)
else: else:
line = f" {cgroup.name:<40} ID: {cgroup.id}" line = f" {cgroup.name:<40} ID: {cgroup.id}"
stdscr.addstr(y, 2, line[:width - 4]) stdscr.addstr(y, 2, line[: width - 4])
# Footer with count # Footer with count
footer = f"Total cgroups: {len(cgroups)}" footer = f"Total cgroups: {len(cgroups)}"
@ -174,21 +173,37 @@ class ContainerMonitorTUI:
# RX graph # RX graph
y += 2 y += 2
rx_label = f"RX: {self._format_bytes(stats.rx_bytes)} ({stats.rx_packets:,} packets)" rx_label = (
f"RX: {self._format_bytes(stats.rx_bytes)} ({stats.rx_packets:,} packets)"
)
stdscr.addstr(y, 2, rx_label) stdscr.addstr(y, 2, rx_label)
if len(history) > 1: if len(history) > 1:
self._draw_bar_graph(stdscr, y + 1, 2, width - 4, 3, self._draw_bar_graph(
[s.rx_bytes for s in history], stdscr,
curses.color_pair(2)) y + 1,
2,
width - 4,
3,
[s.rx_bytes for s in history],
curses.color_pair(2),
)
# TX graph # TX graph
y += 5 y += 5
tx_label = f"TX: {self._format_bytes(stats.tx_bytes)} ({stats.tx_packets:,} packets)" tx_label = (
f"TX: {self._format_bytes(stats.tx_bytes)} ({stats.tx_packets:,} packets)"
)
stdscr.addstr(y, 2, tx_label) stdscr.addstr(y, 2, tx_label)
if len(history) > 1: if len(history) > 1:
self._draw_bar_graph(stdscr, y + 1, 2, width - 4, 3, self._draw_bar_graph(
[s.tx_bytes for s in history], stdscr,
curses.color_pair(3)) y + 1,
2,
width - 4,
3,
[s.tx_bytes for s in history],
curses.color_pair(3),
)
# File I/O graphs # File I/O graphs
y += 5 y += 5
@ -196,21 +211,37 @@ class ContainerMonitorTUI:
# Read graph # Read graph
y += 2 y += 2
read_label = f"READ: {self._format_bytes(stats.read_bytes)} ({stats.read_ops:,} ops)" read_label = (
f"READ: {self._format_bytes(stats.read_bytes)} ({stats.read_ops:,} ops)"
)
stdscr.addstr(y, 2, read_label) stdscr.addstr(y, 2, read_label)
if len(history) > 1: if len(history) > 1:
self._draw_bar_graph(stdscr, y + 1, 2, width - 4, 3, self._draw_bar_graph(
[s.read_bytes for s in history], stdscr,
curses.color_pair(4)) y + 1,
2,
width - 4,
3,
[s.read_bytes for s in history],
curses.color_pair(4),
)
# Write graph # Write graph
y += 5 y += 5
write_label = f"WRITE: {self._format_bytes(stats.write_bytes)} ({stats.write_ops:,} ops)" write_label = (
f"WRITE: {self._format_bytes(stats.write_bytes)} ({stats.write_ops:,} ops)"
)
stdscr.addstr(y, 2, write_label) stdscr.addstr(y, 2, write_label)
if len(history) > 1: if len(history) > 1:
self._draw_bar_graph(stdscr, y + 1, 2, width - 4, 3, self._draw_bar_graph(
[s.write_bytes for s in history], stdscr,
curses.color_pair(5)) y + 1,
2,
width - 4,
3,
[s.write_bytes for s in history],
curses.color_pair(5),
)
def _draw_section_header(self, stdscr, y: int, title: str, color_pair: int): def _draw_section_header(self, stdscr, y: int, title: str, color_pair: int):
"""Draw a section header.""" """Draw a section header."""
@ -220,8 +251,16 @@ class ContainerMonitorTUI:
stdscr.addstr(y, len(title) + 3, "─" * (width - len(title) - 5)) stdscr.addstr(y, len(title) + 3, "─" * (width - len(title) - 5))
stdscr.attroff(curses.color_pair(color_pair) | curses.A_BOLD) stdscr.attroff(curses.color_pair(color_pair) | curses.A_BOLD)
def _draw_bar_graph(self, stdscr, y: int, x: int, width: int, height: int, def _draw_bar_graph(
data: List[float], color_pair: int): self,
stdscr,
y: int,
x: int,
width: int,
height: int,
data: List[float],
color_pair: int,
):
"""Draw a simple bar graph.""" """Draw a simple bar graph."""
if not data or width < 2: if not data or width < 2:
return return
@ -250,25 +289,21 @@ class ContainerMonitorTUI:
else: else:
bar_line += " " bar_line += " "
try: stdscr.attron(color_pair)
stdscr.attron(color_pair) stdscr.addstr(y + row, x, bar_line[:width])
stdscr.addstr(y + row, x, bar_line[:width]) stdscr.attroff(color_pair)
stdscr.attroff(color_pair)
except:
pass # Ignore errors at screen edges
def _format_bytes(self, bytes_val: int) -> str: def _format_bytes(self, bytes_val: float) -> str:
"""Format bytes into human-readable string.""" """Format bytes into human-readable string."""
for unit in ['B', 'KB', 'MB', 'GB', 'TB']: for unit in ["B", "KB", "MB", "GB", "TB"]:
if bytes_val < 1024.0: if bytes_val < 1024.0:
return f"{bytes_val:.2f} {unit}" return f"{bytes_val:.2f} {unit}"
bytes_val /= 1024. bytes_val /= 1024.0
0
return f"{bytes_val:.2f} PB" return f"{bytes_val:.2f} PB"
def _handle_input(self, key: int) -> bool: def _handle_input(self, key: int) -> bool:
"""Handle keyboard input. Returns False to exit.""" """Handle keyboard input. Returns False to exit."""
if key == ord('q') or key == ord('Q'): if key == ord("q") or key == ord("Q"):
return False # Exit return False # Exit
if self.current_screen == "selection": if self.current_screen == "selection":
@ -277,19 +312,19 @@ class ContainerMonitorTUI:
elif key == curses.KEY_DOWN: elif key == curses.KEY_DOWN:
cgroups = self.collector.get_all_cgroups() cgroups = self.collector.get_all_cgroups()
self.selected_index = min(len(cgroups) - 1, self.selected_index + 1) self.selected_index = min(len(cgroups) - 1, self.selected_index + 1)
elif key == ord('\n') or key == curses.KEY_ENTER or key == 10: elif key == ord("\n") or key == curses.KEY_ENTER or key == 10:
# Select cgroup # Select cgroup
cgroups = self.collector.get_all_cgroups() cgroups = self.collector.get_all_cgroups()
if cgroups and 0 <= self.selected_index < len(cgroups): if cgroups and 0 <= self.selected_index < len(cgroups):
cgroups.sort(key=lambda c: c.name) cgroups.sort(key=lambda c: c.name)
self.selected_cgroup = cgroups[self.selected_index].id self.selected_cgroup = cgroups[self.selected_index].id
self.current_screen = "monitoring" self.current_screen = "monitoring"
elif key == ord('r') or key == ord('R'): elif key == ord("r") or key == ord("R"):
# Force refresh cache # Force refresh cache
self.collector._cgroup_cache_time = 0 self.collector._cgroup_cache_time = 0
elif self.current_screen == "monitoring": elif self.current_screen == "monitoring":
if key == 27 or key == ord('b') or key == ord('B'): # ESC or 'b' if key == 27 or key == ord("b") or key == ord("B"): # ESC or 'b'
self.current_screen = "selection" self.current_screen = "selection"
self.selected_cgroup = None self.selected_cgroup = None