#!/usr/bin/python3

"""This program monitors the WOL packets coming into the primary network
interface, then writes to a FIFO that can be read by user programs to
determine if WOL packets have arrived."""

import binascii
import errno
import grp
import os
import multiprocessing
import select
import socket
import sys
import threading
import time

import psutil


SIXBYTES = bytes([255]) * 6


def report(text: str, *args):
    if args:
        text = text % args
    print(text, file=sys.stderr)


def get_macs():
    for interface, snics in psutil.net_if_addrs().items():
        for snic in snics:
            if snic.family == psutil.AF_LINK:
                if snic.address == "00:00:00:00:00:00":
                    continue
                yield snic.address


def wol_receive(
    server: socket.socket,
    queue: multiprocessing.Queue,
):
    clients: list[socket.socket] = []
    report("Started monitoring server socket.")
    while True:
        r, w, e = select.select(
            [server, queue._reader] + clients,  # type: ignore
            [],
            clients,
            None,
        )
        if r:
            if r[0] == server:
                try:
                    client, addr = server.accept()
                except Exception as exc:
                    report("Error accepting connection from client: %s", exc)
                else:
                    report("New client %s.", client)
                    clients.append(client)
            elif r[0] == queue._reader:  # type:ignore
                _ = queue.get()
                _, w2, _ = select.select([], clients, [], 0)
                if w2:
                    report("Notifying %s clients of wakeup.", len(w2))
                for client in w2:
                    try:
                        client.send(b"1")
                    except Exception as exc:
                        report(
                            "Error notifying client %s: %s.  Removing.",
                            client,
                            exc,
                        )
                        client.close()
                        clients.remove(client)
            elif r[0] in clients:
                client = r[0]
                try:
                    c = client.recv(1)
                except Exception as exc:
                    report(
                        "Error reading from client %s: %s.  Removing.",
                        client,
                        exc,
                    )
                else:
                    if len(c) == 0:
                        report(
                            "Client %s sent zero-length message.  Removing.",
                            client,
                        )
                    else:
                        report(
                            "Client %s wrote to us.  Removing.",
                            client,
                        )
                client.close()
                clients.remove(client)
        elif e:
            for client in e:
                report("Client %s had error.  Removing.", client)
                client.close()
                clients.remove(client)


def is_magic_packet(data: bytes) -> bool:
    # Can we find the lead?
    index = data.find(SIXBYTES)
    if index == -1:
        report("Did not find the indicator")
        return False
    index = index + len(SIXBYTES)

    # Can we read all six copies of the MAC?
    first, second, third, fourth, fifth, sixth = (
        data[index : index + 6],  # noqa
        data[index + 6 : index + 12],  # noqa
        data[index + 12 : index + 18],  # noqa
        data[index + 18 : index + 24],  # noqa
        data[index + 24 : index + 30],  # noqa
        data[index + 30 : index + 36],  # noqa
    )

    # Are all six copies the same?
    if (
        first != second
        or second != third
        or third != fourth
        or fourth != fifth
        or fifth != sixth
    ):
        report("Did not find the MACs to be the same")
        return False

    # Does one of the known MAC addresses match the sent data?
    mac = binascii.hexlify(first, ":").decode("ascii")
    my_macs = get_macs()
    for my_mac in my_macs:
        if my_mac == mac:
            return True
    report("Did not find my MAC addresses")
    return False


def wol_monitor():
    os.umask(0o037)
    try:
        os.makedirs("/run/wol-monitor")
    except OSError as e:
        if e.errno != errno.EEXIST:
            raise

    os.chmod("/run/wol-monitor", 0o750)

    if os.path.exists("/run/wol-monitor/wol-bus"):
        os.remove("/run/wol-monitor/wol-bus")

    server = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
    server.bind("/run/wol-monitor/wol-bus")
    os.chmod("/run/wol-monitor/wol-bus", 0o660)

    try:
        group = sys.argv[1]
        try:
            gid = grp.getgrnam(group).gr_gid
        except KeyError:
            print(
                "Error: the UNIX group %r does not exist." % group,
                file=sys.stderr,
            )
            sys.exit(os.EX_USAGE)
        os.chown("/run/wol-monitor/wol-bus", os.getuid(), gid)
        os.chown("/run/wol-monitor", os.getuid(), gid)
    except IndexError:
        pass

    server.listen(2)
    queue = multiprocessing.Queue()
    threading.Thread(target=wol_receive, args=(server, queue)).start()

    publicip = ""
    ports = [9, 40000]
    socks = [
        socket.socket(socket.AF_INET, socket.SOCK_DGRAM, socket.IPPROTO_UDP)
        for _ in ports
    ]
    [s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) for s in socks]
    [s.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1) for s in socks]
    [s.setblocking(False) for s in socks]
    [s.bind((publicip, ports[i])) for i, s in enumerate(socks)]

    report("Started monitoring WOL sockets.")

    ret = 0
    try:
        while True:
            r, w, err = select.select(socks, [], socks, None)
            if err:
                p = ports[socks.index(err[0])]
                report(
                    "Error monitoring WOL socket on port %d.  Exiting.",
                    p,
                )
                ret = 4
                break
            if r:
                data = r[0].recv(9000)
                if not is_magic_packet(data):
                    continue

                p = ports[socks.index(r[0])]
                report(
                    "WOL socket on port %d received wake.  Sending signal.",
                    p,
                )
                queue.put(p)
                time.sleep(10)
                report("Ready for more WOL events.")
                # Now discard whatever may be in the socket, to prevent DoS
                while True:
                    r, _, _ = select.select(socks, [], [], 0)
                    if r:
                        r[0].recv(9000)
                    else:
                        break
    except KeyboardInterrupt:
        pass

    return ret


if __name__ == "__main__":
    sys.exit(wol_monitor())
