PiStatsDisplay/fetchers.py

import os
import time
from datetime import datetime
import datetime as datetimesrc
import psutil
import subprocess
from colors import Colors, getColorsDict
import ipaddress

def pretty_mac(_mac):
    try:
        mac_split = _mac.split(":")
        mac = ""
        for ch in mac_split:
            if len(ch) == 2:
                mac += f"{ch}:"
            if len(ch) == 1:
                mac += f"0{ch}:"
            
        mac = mac[:-1]
        return mac
    except:
        return _mac

class MacLookup:
    prefixes = {}
    #https://github.com/bauerj/mac_vendor_lookup/blob/master/mac_vendor_lookup.py
    def __init__(self):
        self.load_vendors()
        pass

    @staticmethod
    def sanitise(_mac):
        mac = _mac.replace(":", "").replace("-", "").replace(".", "").upper()
        
        try:
            int(mac, 16)
        except ValueError:
            raise Exception("{} contains unexpected character".format(_mac))
        if len(mac) > 12:
            raise Exception("{} is not a valid MAC address (too long)".format(_mac))
        return mac
    
    def lookup(self, mac):
        try:
            mac = self.sanitise(mac)
        except:
            return f"Unknown MAC ({mac})"
        if not self.prefixes:
            self.load_vendors()
        if type(mac) == str:
            mac = mac.encode("utf8")
        try:
            return self.prefixes[mac[:6]].decode("utf8")
        except KeyError:
            return f"Unknown MAC ({mac})"
        
    def load_vendors(self):
        from pathlib import Path
        script_dir = Path(__file__).parent.resolve()
        with open(os.path.join(script_dir, "oui.txt"), "rb") as macs:
            for line in macs.readlines():
                if not line:
                    break
                if b"(base 16)" in line:
                    prefix, vendor = (i.strip() for i in line.split(b"(base 16)", 1))
                    self.prefixes[prefix] = vendor

global maclookup
maclookup = MacLookup()

global firstMacLookup
firstMacLookup = {}

def getAvg():
    data = os.getloadavg()
    return '{0:.2f} {1:.2f} {2:.2f}'.format(*data)

def getDateTime():
    return datetime.now().strftime("%Y-%m-%d %H:%M:%S")

def getCpuPercent(colors = True):
    if colors:
        u = psutil.cpu_percent(interval=1)
        color = ""
        if u > 0:
            color = Colors.GREEN
        if u > 33:
            color = Colors.YELLOW
        if u > 66:
            color = Colors.RED
        if u > 99:
            u = 99
        return "{0}{1:3}%{2}".format(color, u, Colors.END)
    else:
        return "{0:3}%".format(psutil.cpu_percent(interval=1))

def getMemPercent(colors = True):
    if colors:
        u = psutil.virtual_memory().percent
        color = ""
        if u > 0:
            color = Colors.GREEN
        if u > 33:
            color = Colors.YELLOW
        if u > 66:
            color = Colors.RED
        if u > 99:
            u = 99
        return "{0}{1:3}%{2}".format(color, u, Colors.END)
    else:
        return "{0:3}%".format(psutil.virtual_memory().percent)

def getRowSeparator():
    return Colors.NEGATIVE + "".join([" " for i in range(0, 60)]) + Colors.END

def checkIpInPool(ip, pool):
    if not ip or not pool:
        return True
    return ipaddress.ip_address(ip) in ipaddress.ip_network(pool)

def pingHostViaNMap(ip_pool = None):
    if ip_pool is None:
        return ""
    subprocess.check_output(f"nmap -sn {ip_pool}", shell=True, text=True)
    return getDateTime()

def getArp(ip_pool = None):
    #pingHostViaNMap(ip_pool)

    global maclookup
    if maclookup is None:
        maclookup = MacLookup()

    global firstMacLookup
    #linux | ? (192.168.1.136) at a8:7e:ea:64:85:de [ether] on eth0
    #mac | mdns.mcast.net (224.0.0.251) at 1:0:5e:0:0:fb on en0 ifscope permanent [ethernet]
    arp_res = subprocess.check_output("arp -a", shell=True, text=True)
    result = []

    found_macs = []
    for line in arp_res.split("\n"):
        if line:
            host, ip, at, mac, *other = line.split(" ")
            ip = ip[1:-1]
            if checkIpInPool(ip, ip_pool):
                pass
            else:
                continue

            if mac == "(incomplete)":
                continue

            if mac == "FF:FF:FF:FF:FF:FF".lower():
                continue

            mac = pretty_mac(mac)

            if mac not in firstMacLookup:
                firstMacLookup[mac] = int(time.time())
            
            found_macs.append(mac)
            delta = int(time.time()) - firstMacLookup[mac]
            result.append({
                "host": host, 
                "ip": ip, 
                "mac": mac, 
                "mac_name": maclookup.lookup(mac), 
                "delta": delta,
                "mac_first_lookup_utime": str(datetimesrc.timedelta(seconds=delta))
                })

    unused_macs = [mac for mac in list(firstMacLookup.keys()) if mac not in found_macs]
    for mac in unused_macs:
        del firstMacLookup[mac]

    txt = ""
    exists = []
    for host in result:
        host.update(getColorsDict())
        if host["ip"] not in exists:
            exists.append(host["ip"])
            host["time_color"] = Colors.RED

            if host["delta"] > 60 * 60:
                host["time_color"] = Colors.YELLOW

            if host["delta"] > 60 * 60 * 12:
                host["time_color"] = Colors.GREEN

            mac_name_len = len(host["mac_name"])
            host["mac_name"] = host["mac_name"][:31] + ("..." if mac_name_len > 31 else "")

            txt += "{time_color}{mac_first_lookup_utime}{END} - {BLUE}{ip}{END} - {mac_name}\n".format(**host)
    return txt[:-1]

if __name__ == "__main__":
    print(getArp())