#!/usr/bin/env python3
"""
SSD1306 128x64 OLED — Five-page Network Monitor with custom flip times

Page 0 (Start): 
  tcpdump (centered)        [8s]
  BPF: <expr|none> (centered)
  [blank]
  Chappell University (centered)

Page 1 (Identity): host/IP/MAC/iface
Page 2 (Status): totals, uptime, date, time
Page 3 (Protocols): totals + TCP/UDP/ICMP
Page 4 (PPS): packets/sec IN/OUT (ints)
"""

import argparse, subprocess, re, time, sys, socket, datetime, select
from pathlib import Path
from PIL import Image, ImageDraw, ImageFont

import board, busio
from adafruit_ssd1306 import SSD1306_I2C

MAC_RE = re.compile(
    r'([0-9A-Fa-f]{2}(?::[0-9A-Fa-f]{2}){5})\s*>\s*([0-9A-Fa-f]{2}(?::[0-9A-Fa-f]{2}){5})'
)
PROTO_PORT_LINE = re.compile(r'\bIP6?\b .*?\.\d+\s*>\s*.*?\.\d+:')
def classify_protocol(line: str) -> str | None:
    L = line
    if "ICMP6" in L or "icmp6" in L or "ICMPv6" in L: return "icmp"
    if re.search(r'\bICMP\b', L, re.I): return "icmp"
    if re.search(r'\bUDP\b', L, re.I): return "udp"
    if "Flags [" in L: return "tcp"
    if re.search(r'\bTCP\b', L, re.I): return "tcp"
    if PROTO_PORT_LINE.search(L):
        if re.search(r'\bUDP\b', L, re.I): return "udp"
        return "tcp"
    return None

def read_local_mac(iface: str) -> str:
    return Path(f"/sys/class/net/{iface}/address").read_text().strip().lower()
def read_local_ip(iface: str) -> str:
    try:
        out = subprocess.check_output(["ip","-4","-o","addr","show","dev",iface], text=True)
        m = re.search(r"\binet\s+(\d+\.\d+\.\d+\.\d+)/\d+\b", out)
        return m.group(1) if m else "No IP"
    except: return "No IP"
def hostname() -> str: return socket.gethostname()

def init_display(width, height, addr):
    i2c = busio.I2C(board.SCL, board.SDA)
    disp = SSD1306_I2C(width, height, i2c, addr=addr)
    disp.fill(0); disp.show()
    return disp
def try_load_font(path, size):
    if path:
        try: return ImageFont.truetype(path, size)
        except: pass
    try: return ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", size)
    except: return ImageFont.load_default()
def line_height(font):
    try: a,d = font.getmetrics(); return max(8,a+d)
    except: b = font.getbbox("Ag"); return max(8,b[3]-b[1])
def fit_font(fontpath, start_size, lines_needed, max_height):
    size = start_size
    while size>=7:
        f=try_load_font(fontpath,size); lh=line_height(f)
        if lh*lines_needed<=max_height: return f,lh
        size-=1
    f=try_load_font(fontpath,7); return f,line_height(f)
def truncate(draw,text,font,max_w):
    if draw.textlength(text,font=font)<=max_w: return text
    ell="…"; target=max_w-draw.textlength(ell,font=font); out=""
    for ch in text:
        if draw.textlength(out+ch,font=font)>target: break
        out+=ch
    return out+ell
def center_text(draw,text,font,w,y):
    tw=draw.textlength(text,font=font); x=max(0,(w-tw)//2)
    draw.text((x,y),text,font=font,fill=255)

def build_tcpdump_cmd(iface,bpf):
    return ["tcpdump","-i",iface,"-n","-e","-l"]+[bpf if bpf else "(ip or ip6)"]

# --- pages ---
def draw_page0_start(disp,font,lh,bpf):
    img=Image.new("1",(disp.width,disp.height)); d=ImageDraw.Draw(img); w,y=disp.width,0
    center_text(d,"tcpdump",font,w,y); y+=lh
    bfps=f"BPF: {bpf if bpf else 'none'}"
    center_text(d,bfps,font,w,y); y+=lh
    y+=lh  # blank line
    center_text(d,"Chappell University",font,w,y)
    disp.image(img); disp.show()
def draw_page1_identity(disp,font,lh,iface,bpf,host,ip,mac):
    img=Image.new("1",(disp.width,disp.height)); d=ImageDraw.Draw(img); max_w,y=disp.width,0
    lines=(f"Host: {host}",f"IP:   {ip}",f"MAC:  {mac}",f"IF: {iface}  BPF: {bpf if bpf else 'none'}")
    for txt in lines: d.text((0,y),truncate(d,txt,font,max_w),font=font,fill=255); y+=lh
    disp.image(img); disp.show()
def format_hms(seconds):
    seconds=int(seconds); h=seconds//3600; m=(seconds%3600)//60; s=seconds%60
    return f"{h:02d}:{m:02d}:{s:02d}"
def draw_page2_status(disp,font,lh,in_cnt,out_cnt,start_ts):
    img=Image.new("1",(disp.width,disp.height)); d=ImageDraw.Draw(img); max_w,y=disp.width,0
    nowdt=datetime.datetime.now()
    lines=(f"Totals: IN {in_cnt}  OUT {out_cnt}",
           f"Uptime: {format_hms(time.time()-start_ts)}",
           f"Date:   {nowdt.strftime('%Y-%m-%d')}",
           f"Time:   {nowdt.strftime('%H:%M:%S')}")
    for txt in lines: d.text((0,y),truncate(d,txt,font,max_w),font=font,fill=255); y+=lh
    disp.image(img); disp.show()
def draw_page3_protocols(disp,font,lh,tcp_cnt,udp_cnt,icmp_cnt):
    img=Image.new("1",(disp.width,disp.height)); d=ImageDraw.Draw(img); max_w,y=disp.width,0
    total=tcp_cnt+udp_cnt+icmp_cnt
    lines=(f"Total pkts: {total}",f"TCP: {tcp_cnt}",f"UDP: {udp_cnt}",f"ICMP: {icmp_cnt}")
    for txt in lines: d.text((0,y),truncate(d,txt,font,max_w),font=font,fill=255); y+=lh
    disp.image(img); disp.show()
def draw_page4_pps(disp,font,lh,in_pps,out_pps):
    img=Image.new("1",(disp.width,disp.height)); d=ImageDraw.Draw(img); max_w,y=disp.width,0
    in_i=int(round(in_pps)); out_i=int(round(out_pps))
    # If you want OUT on the last line with a blank third line, use: ("Packets per Second", f"IN: {in_i}", "", f"OUT: {out_i}")
    lines=("Packets per Second",f"IN: {in_i}",f"OUT: {out_i}","")
    for txt in lines:
        if txt: d.text((0,y),truncate(d,txt,font,max_w),font=font,fill=255)
        y+=lh
    disp.image(img); disp.show()

def main():
    ap=argparse.ArgumentParser()
    ap.add_argument("--iface",default="wlan0"); ap.add_argument("--bpf",default=None)
    ap.add_argument("--addr",type=lambda x:int(x,0),default=0x3C)
    ap.add_argument("--width",type=int,default=128); ap.add_argument("--height",type=int,default=64)
    ap.add_argument("--fontsize",type=int,default=12); ap.add_argument("--fontpath",default=None)
    ap.add_argument("--update",type=float,default=0.25)
    ap.add_argument("--page",type=int,choices=[0,1,2,3,4],default=None)
    args=ap.parse_args()

    local_mac=read_local_mac(args.iface); local_ip=read_local_ip(args.iface); host=hostname()
    disp=init_display(args.width,args.height,args.addr)
    font,lh=fit_font(args.fontpath,args.fontsize,4,args.height)

    in_cnt=out_cnt=tcp_cnt=udp_cnt=icmp_cnt=0
    last_ts=time.time(); last_in=last_out=0; in_pps=out_pps=0.0; start_ts=time.time()

    draw_page0_start(disp,font,lh,args.bpf)
    cmd=build_tcpdump_cmd(args.iface,args.bpf)
    try:
        proc=subprocess.Popen(cmd,stdout=subprocess.PIPE,stderr=subprocess.DEVNULL,text=True,bufsize=1)
    except Exception as e:
        sys.exit(str(e))

    cur_page=0
    # First flip after 8s on start page; then 5s elsewhere
    next_flip=time.time()+ (8 if args.page is None else 10**9)

    try:
        # Non-blocking loop: flip/draw even when no packets match BPF
        while True:
            # Wait up to args.update seconds for a tcpdump line
            rlist, _, _ = select.select([proc.stdout], [], [], args.update)
            if rlist:
                line = proc.stdout.readline()
                if not line:
                    break  # tcpdump ended

                # IN/OUT (MAC-based)
                m=MAC_RE.search(line)
                if m:
                    src,dst=m.group(1).lower(),m.group(2).lower()
                    if src==local_mac: out_cnt+=1
                    elif dst==local_mac: in_cnt+=1

                # Protocol counts
                proto=classify_protocol(line)
                if proto=="icmp": icmp_cnt+=1
                elif proto=="udp": udp_cnt+=1
                elif proto=="tcp": tcp_cnt+=1

            # Tick/update even if no new line arrived
            now=time.time()
            if now-last_ts>=args.update:
                dt=now-last_ts; din=in_cnt-last_in; dout=out_cnt-last_out
                in_pps=din/dt if dt>0 else 0.0
                out_pps=dout/dt if dt>0 else 0.0
                last_ts,last_in,last_out=now,in_cnt,out_cnt

                # Page selection / flip
                if args.page is not None:
                    cur_page=args.page
                else:
                    if now>=next_flip:
                        cur_page=0 if cur_page==4 else cur_page+1
                        next_flip=now+(8 if cur_page==0 else 5)

                # Draw current page
                if cur_page==0: draw_page0_start(disp,font,lh,args.bpf)
                elif cur_page==1: draw_page1_identity(disp,font,lh,args.iface,args.bpf,host,local_ip,local_mac)
                elif cur_page==2: draw_page2_status(disp,font,lh,in_cnt,out_cnt,start_ts)
                elif cur_page==3: draw_page3_protocols(disp,font,lh,tcp_cnt,udp_cnt,icmp_cnt)
                else: draw_page4_pps(disp,font,lh,in_pps,out_pps)

    except KeyboardInterrupt:
        pass
    finally:
        try: proc.terminate()
        except: pass
        disp.fill(0); disp.show()

if __name__=="__main__": main()