printer_levels.py

#!/usr/bin/env python3
import argparse
import ipaddress
import socket
import sys
from concurrent.futures import ThreadPoolExecutor, as_completed
from dataclasses import dataclass
from typing import Dict, List, Optional, Tuple

import psutil

from pysnmp.hlapi.asyncio import (
    SnmpEngine,
    CommunityData,
    UdpTransportTarget,
    ContextData,
    ObjectType,
    ObjectIdentity,
    getCmd,
    nextCmd,
)


# -----------------------------
# Standard Printer-MIB OIDs
# -----------------------------
# prtGeneralPrinterName.0
OID_PRINTER_NAME = "1.3.6.1.2.1.43.5.1.1.16.1"

# Supplies table (indexed)
# prtMarkerSuppliesDescription
OID_SUP_DESC = "1.3.6.1.2.1.43.11.1.1.6.1"
# prtMarkerSuppliesLevel
OID_SUP_LEVEL = "1.3.6.1.2.1.43.11.1.1.9.1"
# prtMarkerSuppliesMaxCapacity
OID_SUP_MAX = "1.3.6.1.2.1.43.11.1.1.8.1"
# prtMarkerSuppliesClass (1=other,3=supplyThatIsConsumed,4=receptacleThatIsFilled, etc.)
OID_SUP_CLASS = "1.3.6.1.2.1.43.11.1.1.4.1"
# prtMarkerSuppliesType (varies)
OID_SUP_TYPE = "1.3.6.1.2.1.43.11.1.1.5.1"

# Some printers report "unknown" as -3 for level; some report -2 for "not available"
UNKNOWN_LEVEL_VALUES = {-3, -2}

@dataclass
class Supply:
    index: int
    description: str
    level: Optional[int]
    max_capacity: Optional[int]
    pct: Optional[float]

@dataclass
class PrinterInfo:
    ip: str
    name: str
    supplies: List[Supply]

def _snmp_get(ip: str, community: str, oid: str, timeout: float, retries: int) -> Optional[str]:
    """SNMP GET for a scalar OID. Returns stringified value or None."""
    try:
        iterator = getCmd(
            SnmpEngine(),
            CommunityData(community, mpModel=1),  # SNMPv2c
            UdpTransportTarget((ip, 161), timeout=timeout, retries=retries),
            ContextData(),
            ObjectType(ObjectIdentity(oid)),
        )
        errorIndication, errorStatus, errorIndex, varBinds = next(iterator)
        if errorIndication or errorStatus:
            return None
        for _, val in varBinds:
            return str(val)
        return None
    except Exception:
        return None

def _snmp_walk(ip: str, community: str, base_oid: str, timeout: float, retries: int, max_rows: int = 256) -> List[Tuple[str, str]]:
    """SNMP WALK starting at base_oid. Returns list of (oid, value)."""
    results: List[Tuple[str, str]] = []
    try:
        for (errorIndication, errorStatus, errorIndex, varBinds) in nextCmd(
            SnmpEngine(),
            CommunityData(community, mpModel=1),
            UdpTransportTarget((ip, 161), timeout=timeout, retries=retries),
            ContextData(),
            ObjectType(ObjectIdentity(base_oid)),
            lexicographicMode=False,
        ):
            if errorIndication or errorStatus:
                break
            for oid, val in varBinds:
                oid_str = str(oid)
                if not oid_str.startswith(base_oid + ".") and oid_str != base_oid:
                    return results
                results.append((oid_str, str(val)))
                if len(results) >= max_rows:
                    return results
        return results
    except Exception:
        return results

def _parse_index(oid: str, base: str) -> Optional[int]:
    # base.x where x is the row index
    if not oid.startswith(base + "."):
        return None
    tail = oid[len(base) + 1 :]
    try:
        return int(tail)
    except ValueError:
        return None

def query_printer(ip: str, community: str, timeout: float, retries: int) -> Optional[PrinterInfo]:
    # quick check: try to read printer name; if it fails, assume not a printer / no SNMP
    name = _snmp_get(ip, community, OID_PRINTER_NAME, timeout, retries)
    if not name:
        return None

    # Walk supplies columns
    desc_rows = _snmp_walk(ip, community, OID_SUP_DESC, timeout, retries)
    lvl_rows  = _snmp_walk(ip, community, OID_SUP_LEVEL, timeout, retries)
    max_rows  = _snmp_walk(ip, community, OID_SUP_MAX, timeout, retries)

    desc_by_i: Dict[int, str] = {}
    lvl_by_i: Dict[int, int] = {}
    max_by_i: Dict[int, int] = {}

    for oid, val in desc_rows:
        idx = _parse_index(oid, OID_SUP_DESC)
        if idx is not None:
            desc_by_i[idx] = val

    for oid, val in lvl_rows:
        idx = _parse_index(oid, OID_SUP_LEVEL)
        if idx is not None:
            try:
                lvl_by_i[idx] = int(val)
            except ValueError:
                pass

    for oid, val in max_rows:
        idx = _parse_index(oid, OID_SUP_MAX)
        if idx is not None:
            try:
                max_by_i[idx] = int(val)
            except ValueError:
                pass

    supplies: List[Supply] = []
    all_indexes = sorted(set(desc_by_i.keys()) | set(lvl_by_i.keys()) | set(max_by_i.keys()))
    for i in all_indexes:
        desc = desc_by_i.get(i, f"Supply {i}")
        level_raw = lvl_by_i.get(i, None)
        max_raw = max_by_i.get(i, None)

        level: Optional[int] = None
        max_capacity: Optional[int] = None
        pct: Optional[float] = None

        if level_raw is not None and level_raw not in UNKNOWN_LEVEL_VALUES:
            level = level_raw
        if max_raw is not None and max_raw > 0:
            max_capacity = max_raw

        if level is not None and max_capacity is not None:
            # Clamp to sane range
            pct = max(0.0, min(100.0, (level / max_capacity) * 100.0))

        supplies.append(Supply(index=i, description=desc, level=level, max_capacity=max_capacity, pct=pct))

    return PrinterInfo(ip=ip, name=name, supplies=supplies)

def udp_161_open(ip: str, timeout: float) -> bool:
    """Best-effort probe: send a dummy UDP packet to 161 and rely on SNMP GET later.
    UDP 'open' isn't reliably detectable; we keep this lightweight and just attempt SNMP in query_printer.
    """
    return True

def get_local_cidrs() -> List[ipaddress.IPv4Network]:
    cidrs: List[ipaddress.IPv4Network] = []
    for iface, addrs in psutil.net_if_addrs().items():
        for a in addrs:
            if a.family == socket.AF_INET and a.address and a.netmask:
                ip = ipaddress.IPv4Address(a.address)
                if ip.is_loopback or ip.is_link_local:
                    continue
                net = ipaddress.IPv4Network(f"{a.address}/{a.netmask}", strict=False)
                if net.prefixlen <= 16:
                    # avoid scanning huge ranges by default
                    continue
                cidrs.append(net)
    # de-dupe
    uniq = []
    seen = set()
    for n in cidrs:
        if str(n) not in seen:
            uniq.append(n)
            seen.add(str(n))
    return uniq

def format_supply(s: Supply) -> str:
    if s.pct is not None:
        return f"{s.description}: {s.pct:.0f}% ({s.level}/{s.max_capacity})"
    if s.level is not None and s.max_capacity is None:
        return f"{s.description}: level={s.level}"
    if s.level is None and s.max_capacity is not None:
        return f"{s.description}: max={s.max_capacity} (level unknown)"
    return f"{s.description}: unknown"

def main():
    ap = argparse.ArgumentParser(description="Discover SNMP printers and query toner/ink levels (Printer-MIB).")
    ap.add_argument("--cidr", action="append", help="CIDR to scan (can repeat). Example: 192.168.1.0/24")
    ap.add_argument("--community", default="public", help="SNMP community string (default: public)")
    ap.add_argument("--timeout", type=float, default=0.8, help="SNMP timeout seconds (default: 0.8)")
    ap.add_argument("--retries", type=int, default=0, help="SNMP retries (default: 0)")
    ap.add_argument("--workers", type=int, default=128, help="Parallel workers (default: 128)")
    ap.add_argument("--max-hosts", type=int, default=4096, help="Safety cap for total hosts scanned (default: 4096)")
    args = ap.parse_args()

    cidrs: List[ipaddress.IPv4Network] = []
    if args.cidr:
        for c in args.cidr:
            cidrs.append(ipaddress.IPv4Network(c, strict=False))
    else:
        cidrs = get_local_cidrs()

    if not cidrs:
        print("No suitable local /24-/30 networks detected. Provide --cidr 192.168.1.0/24")
        sys.exit(2)

    targets: List[str] = []
    for net in cidrs:
        for ip in net.hosts():
            targets.append(str(ip))

    if len(targets) > args.max_hosts:
        print(f"Refusing to scan {len(targets)} hosts (cap {args.max_hosts}). Provide smaller --cidr or raise --max-hosts.")
        sys.exit(2)

    printers: List[PrinterInfo] = []

    with ThreadPoolExecutor(max_workers=args.workers) as ex:
        futures = [ex.submit(query_printer, ip, args.community, args.timeout, args.retries) for ip in targets]
        for fut in as_completed(futures):
            p = fut.result()
            if p:
                printers.append(p)

    printers.sort(key=lambda x: (x.name.lower(), x.ip))

    if not printers:
        print("No SNMP printers found. Try:")
        print("  - different --community (many are 'public', some are custom)")
        print("  - enable SNMP on the printers")
        print("  - scan the right subnet with --cidr")
        sys.exit(0)

    for p in printers:
        print("=" * 72)
        print(f"{p.name}  ({p.ip})")
        if not p.supplies:
            print("  Supplies: (no Printer-MIB supplies reported)")
        else:
            for s in p.supplies:
                print("  - " + format_supply(s))

if __name__ == "__main__":
    main()