#!/usr/bin/env python3
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
"""
Obtain credentials for Actions runner and co-operate with auto scaling group

The purpose of this script is to enable the self-hosted runners to operate in
an AutoScaling environment (without needing admin permissions on the GitHub
repo to create and delete runners.)

The order of operations is:

1. Obtain credentials

   We have pre-created a number of credentials and stored them in Amazon KMS.
   This script makes use of dynamodb to obtain an exclusive lock on a set of
   credentials.

   We need the "locking" as if you use credentials that are already
   in use the new runner process will wait (but never error) until they are not
   in use.

2. Complete the ASG lifecycle action so the instance is marked as InService

   This might not be strictly necessary, we don't want the instance to be "in
   service" until the runner has started.

3. Emit metric saying whether instance is running a job or not

   This is used to drive the scale-in CloudWatch alarm

4. Monitor for the runner starting jobs, and protecting the instance from Scale-In when it is

   Since we are running in an autoscaling group we can't dictate which instance
   AWS choses to terminate, so we instead have to set scale-in protection when a job is running.

   The way we watch for jobs being executed is using the Netlink Process
   Connector, which is a datagram socket that a (root) process can open to the
   kernel, to receive push events for whenever a process starts or stops.

   There are more events than that send, and to limit it to the only ones we
   care about we use a BPF filter to drop everything else.

   Since it is a datagram socket it is possible we might miss a notification, so
   we also periodically check if the process is still alive

5. Watch for ASG instance state changing to Terminating:Wait

   When the ASG wants to terminate the instance, we have it configured to put
   the instance in to a "requested" state -- this is to avoid a race condition
   where the instance _isn't_ running a job (so isn't protected from scale in),
   gets set to Terminating, but before AWS shuts down the machine the runner
   process picks up and starts a Job, which leads to the job failing with "The
   self-hosted runner: Airflow Runner $N lost communication with the server".

   When we notice being in this state, we _gracefully_ shut down the runner
   (letting it complete any job it might have), stop it from restarting, and
   then allow the termination lifecycle to continue

"""
import ctypes
import datetime
import enum
import errno
import json
import logging
import os
import random
import selectors
import shutil
import signal
import socket
from subprocess import check_call
from typing import Callable, List, Tuple, Union

import boto3
import click
import psutil
from python_dynamodb_lock.python_dynamodb_lock import DynamoDBLockClient, DynamoDBLockError
from tenacity import before_sleep_log, retry, stop_after_delay, wait_exponential

logging.basicConfig(level=logging.INFO)
log = logging.getLogger(__name__)
log.setLevel(logging.DEBUG)

logging.getLogger('python_dynamodb_lock').setLevel(logging.WARNING)


TABLE_NAME = os.getenv('COUNTER_TABLE', 'GithubRunnerQueue')


@click.command()
@click.option('--repo', default='apache/airflow')
@click.option('--user', default='runner')
@click.option(
    '--output-folder',
    help="Folder to write credentials to. Default of ~runner/actions-runner",
    default='~runner/actions-runner',
)
def main(repo, output_folder, user):
    global INSTANCE_ID
    # Notify the ASG LifeCycle hook that we are now In Service and ready to
    # process requests/safe to be shut down

    # Fetch current instance ID from where cloutinit writes it to
    if not INSTANCE_ID:
        with open('/var/lib/cloud/data/instance-id') as fh:
            INSTANCE_ID = fh.readline().strip()

    log.info("Starting on %s...", INSTANCE_ID)

    output_folder = os.path.expanduser(output_folder)

    short_time = datetime.timedelta(microseconds=1)

    dynamodb = boto3.resource('dynamodb')
    client = DynamoDBLockClient(
        dynamodb,
        table_name='GitHubRunnerLocks',
        expiry_period=datetime.timedelta(0, 300),
        heartbeat_period=datetime.timedelta(seconds=10),
    )

    # Just keep trying until we get some credentials.
    while True:
        # Have each runner try to get a credential in a random order.
        possibles = get_possible_credentials(repo)
        random.shuffle(possibles)

        log.info("Trying to get a set of credentials in this order: %r", possibles)

        notify = get_sd_notify_func()

        for index in possibles:
            try:
                lock = client.acquire_lock(
                    f'{repo}/{index}',
                    retry_period=short_time,
                    retry_timeout=short_time,
                    raise_context_exception=True,
                )
            except DynamoDBLockError as e:
                log.info("Could not lock %s (%s)", index, e)
                continue

            with lock:
                log.info("Obtained lock on %s", index)
                write_credentials_to_files(repo, index, output_folder, user)
                merge_in_settings(repo, output_folder)
                notify(f"STATUS=Obtained lock on {index}")

                if get_lifecycle_state() == "Pending:Wait":
                    complete_asg_lifecycle_hook()

                notify("READY=1")
                log.info("Watching for Runner.Worker processes")
                ProcessWatcher().run()

            client.close()

            exit()


def get_sd_notify_func() -> Callable[[str], None]:
    # http://www.freedesktop.org/software/systemd/man/sd_notify.html
    addr = os.getenv('NOTIFY_SOCKET')
    if not addr:
        return lambda status: None

    sock = socket.socket(socket.AF_UNIX, socket.SOCK_DGRAM)
    if addr[0] == '@':
        addr = '\0' + addr[1:]
    sock.connect(addr)

    def notify(status: str):
        sock.sendall(status.encode('utf-8'))

    return notify


def write_credentials_to_files(
    repo: str, index: str, out_folder: str = '~runner/actions-runner', user: str = 'runner'
):
    param_path = os.path.join('/runners/', repo, index)

    resp = boto3.client("ssm").get_parameters_by_path(Path=param_path, Recursive=False, WithDecryption=True)

    param_to_file = {
        'config': '.runner',
        'credentials': '.credentials',
        'rsaparams': '.credentials_rsaparams',
    }

    for param in resp['Parameters']:
        # "/runners/apache/airflow/config" -> "config"
        name = os.path.basename(param['Name'])
        filename = param_to_file.get(name, None)
        if filename is None:
            log.info("Unknown Parameter from SSM: %r", param['Name'])
            continue
        log.info("Writing %r to %r", param['Name'], filename)
        with open(os.path.join(out_folder, filename), "w") as fh:
            fh.write(param['Value'])
            shutil.chown(fh.name, user)
            os.chmod(fh.name, 0o600)
        del param_to_file[name]
    if param_to_file:
        raise RuntimeError(f"Missing expected params: {list(param_to_file.keys())}")


def merge_in_settings(repo: str, out_folder: str) -> None:
    client = boto3.client('ssm')

    param_path = os.path.join('/runners/', repo, 'configOverlay')
    log.info("Loading config overlay from %s", param_path)

    try:

        resp = client.get_parameter(Name=param_path, WithDecryption=True)
    except client.exceptions.ParameterNotFound:
        log.debug("Failed to load config overlay", exc_info=True)
        return

    try:
        overlay = json.loads(resp['Parameter']['Value'])
    except ValueError:
        log.debug("Failed to parse config overlay", exc_info=True)
        return

    with open(os.path.join(out_folder, ".runner"), "r+") as fh:
        settings = json.load(fh)

        for key, val in overlay.items():
            settings[key] = val

        fh.seek(0, os.SEEK_SET)
        os.ftruncate(fh.fileno(), 0)
        json.dump(settings, fh, indent=2)


def get_possible_credentials(repo: str) -> List[str]:
    client = boto3.client("ssm")
    paginator = client.get_paginator("describe_parameters")

    path = os.path.join('/runners/', repo, '')
    baked_path = os.path.join(path, 'runnersList')

    # Pre-compute the list, to avoid making lots of requests and getting throttled by SSM API in case of
    # thundering herd
    try:
        log.info("Using pre-computed credentials indexes from %s", baked_path)
        resp = client.get_parameter(Name=baked_path)
        return resp['Parameter']['Value'].split(',')
    except client.exceptions.ParameterNotFound:
        pass

    log.info("Looking at %s for possible credentials", path)

    pages = paginator.paginate(
        ParameterFilters=[{"Key": "Path", "Option": "Recursive", "Values": [path]}],
        PaginationConfig={
            "PageSize": 50,
        },
    )

    seen = set()

    for i, page in enumerate(pages):
        log.info("Page %d", i)
        for param in page['Parameters']:
            name = param['Name']
            log.info("%s", name)

            # '/runners/x/1/config' -> '1/config',
            # '/runners/x/y/1/config' -> 'y/1/config',
            local_name = name[len(path) :]

            try:
                # '1/config' -> '1'
                index, _ = local_name.split('/')
            except ValueError:
                # Ignore any 'x/y' when we asked for 'x'. There should only be an index and a filename
                log.debug("Ignoring nested path %s", name)
                continue

            try:
                # Check it's a number, but keep variable as string
                int(index)
            except ValueError:
                log.debug("Ignoring non-numeric index %s", name)
                continue

            index = os.path.basename(os.path.dirname(name))
            seen.add(index)

    if not seen:
        raise RuntimeError(f'No credentials found in SSM ParameterStore for {repo!r}')

    try:
        resp = client.put_parameter(
            Name=baked_path, Type='StringList', Value=','.join(list(seen)), Overwrite=False
        )
        log.info("Stored pre-computed credentials indexes at %s", baked_path)
    except client.exceptions.ParameterAlreadyExists:
        # Race, we lost, never mind!
        pass

    return list(seen)


OWN_ASG = None
INSTANCE_ID = None


def get_lifecycle_state() -> str:
    global INSTANCE_ID, OWN_ASG

    if not INSTANCE_ID:
        with open('/var/lib/cloud/data/instance-id') as fh:
            INSTANCE_ID = fh.readline().strip()

    asg_client = boto3.client('autoscaling')

    try:
        instances = asg_client.describe_auto_scaling_instances(
            InstanceIds=[INSTANCE_ID],
        )['AutoScalingInstances']
    except asg_client.exceptions.ClientError:
        return "UNKNOWN"

    if len(instances) != 1:
        return "UNKNOWN"

    details = instances[0]

    if not OWN_ASG:
        OWN_ASG = details['AutoScalingGroupName']

    return details['LifecycleState']


def complete_asg_lifecycle_hook(hook_name='WaitForInstanceReportReady', retry=False):
    global OWN_ASG, INSTANCE_ID
    # Notify the ASG LifeCycle hook that we are now InService and ready to
    # process requests/safe to be shut down

    asg_client = boto3.client('autoscaling')

    try:
        asg_client.complete_lifecycle_action(
            AutoScalingGroupName=OWN_ASG,
            InstanceId=INSTANCE_ID,
            LifecycleHookName=hook_name,
            LifecycleActionResult='CONTINUE',
        )
        log.info("LifeCycle hook %s set to CONTINUE instance=%s", hook_name, INSTANCE_ID)
    except asg_client.exceptions.ClientError as e:
        # If the script fails for whatever reason and we re-run, the lifecycle hook may have already be
        # completed, so this would fail. That is not an error

        # We don't want the stacktrace here, just the message
        log.warning("Failed to complete lifecycle hook %s: %s", hook_name, str(e))
        pass


# Constants and types from
# https://github.com/torvalds/linux/blob/fcadab740480e0e0e9fa9bd272acd409884d431a/include/uapi/linux/cn_proc.h
class NlMsgFlag(enum.IntEnum):
    NoOp = 1
    Error = 2
    Done = 3
    Overrun = 4


class NLMsgHdr(ctypes.Structure):
    """Netlink Message Header"""

    _fields_ = [
        ("len", ctypes.c_uint32),
        ("type", ctypes.c_uint16),
        ("flags", ctypes.c_uint16),
        ("seq", ctypes.c_uint32),
        ("pid", ctypes.c_uint32),
    ]


class ProcConnectorOp(enum.IntEnum):
    MCAST_LISTEN = 1
    MCAST_IGNORE = 2


class cn_msg(ctypes.Structure):
    """Linux kernel Connector message"""

    CN_IDX_PROC = 1
    CN_VAL_PROC = 1

    _fields_ = [
        ("cb_id_idx", ctypes.c_uint32),
        ("cb_id_val", ctypes.c_uint32),
        ("seq", ctypes.c_uint32),
        ("ack", ctypes.c_uint32),
        ("len", ctypes.c_uint16),
        ("flags", ctypes.c_uint16),
    ]

    def __init__(self, header, data, **kwargs):
        super().__init__(**kwargs)
        self.header = header
        self.len = ctypes.sizeof(data)
        self.data = data
        self.header.len = ctypes.sizeof(header) + ctypes.sizeof(self) + self.len

    def to_bytes(self):
        return bytes(self.header) + bytes(self) + bytes(self.data)  # type: ignore


class ProcEventWhat(enum.IntFlag):
    NONE = 0x0
    FORK = 0x1
    EXEC = 0x2
    UID = 0x4
    GID = 0x40
    SID = 0x80
    PTRACE = 0x0000010
    COMM = 0x0000020
    COREDUMP = 0x40000000
    EXIT = 0x80000000


class proc_event(ctypes.Structure):
    """Base proc_event field"""

    _fields_ = [
        ("what", ctypes.c_uint32),
        ("cpu", ctypes.c_uint32),
        ("timestamp", ctypes.c_uint64),  # Number of nano seconds since system boot
    ]

    @classmethod
    def from_netlink_packet(
        cls, data
    ) -> Tuple["proc_event", Union[None, "exec_proc_event", "exit_proc_event"]]:
        """
        Parse the netlink packet in to a
        """
        # Netlink message header (struct nlmsghdr)
        header = NLMsgHdr.from_buffer_copy(data)
        data = data[ctypes.sizeof(header) :]

        # We already checked/filtered on header.type == NlMsgFlag.Done

        # Connector message header (struct cn_msg)
        connector_msg = cn_msg.from_buffer_copy(data)

        # Ignore messages from other Netlink connector types: done via BPF

        data = data[ctypes.sizeof(connector_msg) :]

        event = proc_event.from_buffer_copy(data)
        data = data[ctypes.sizeof(event) :]
        event.what = ProcEventWhat(event.what)

        if event.what == ProcEventWhat.EXEC:
            return event, exec_proc_event.from_buffer_copy(data)
        elif event.what == ProcEventWhat.EXIT:
            return event, exit_proc_event.from_buffer_copy(data)
        return event, None


class exec_proc_event(ctypes.Structure):
    _fields_ = [
        ("pid", ctypes.c_int32),
        ("tid", ctypes.c_int32),
    ]


class exit_proc_event(ctypes.Structure):
    _fields_ = [
        ("pid", ctypes.c_int32),
        ("tid", ctypes.c_int32),
        ("exit_code", ctypes.c_int32),
        ("signal", ctypes.c_int32),
    ]


class ProcessWatcher:
    interesting_processes = {}

    protected = None
    in_termating_lifecycle = False

    def run(self):
        # Create a signal pipe that we can poll on
        sig_read, sig_write = socket.socketpair()

        sel = selectors.DefaultSelector()

        def sig_handler(signal, frame):
            # no-op
            ...

        sig_read.setblocking(False)
        sig_write.setblocking(False)
        sel.register(sig_read, selectors.EVENT_READ, None)

        proc_socket = self.open_proc_connector_socket()
        proc_socket.setblocking(False)

        signal.signal(signal.SIGINT, sig_handler)
        signal.signal(signal.SIGALRM, sig_handler)
        signal.setitimer(signal.ITIMER_REAL, 30, 30.0)
        signal.set_wakeup_fd(sig_write.fileno(), warn_on_full_buffer=False)

        sel.register(proc_socket, selectors.EVENT_READ, self.handle_proc_event)

        self.pgrep()

        try:
            while True:
                for key, mask in sel.select():

                    if key.fileobj == sig_read:
                        sig = signal.Signals(key.fileobj.recv(1)[0])  # type: ignore
                        if sig == signal.SIGALRM:
                            self.check_still_alive()
                            continue
                        else:
                            log.info(f"Got {sig.name}, exiting")
                            return
                    callback = key.data
                    callback(key.fileobj, mask)
        finally:
            # Disable the timers for any cleanup code to run
            signal.setitimer(signal.ITIMER_REAL, 0)
            signal.set_wakeup_fd(-1)

    def pgrep(self):
        """Check for any interesting processes we might have missed."""
        listener_found = False

        for proc in psutil.process_iter(['name', 'cmdline']):
            if proc.name() == "Runner.Worker" and proc.pid not in self.interesting_processes:
                log.info(
                    "Found existing interesting processes, protecting from scale in %d: %s",
                    proc.pid,
                    proc.cmdline(),
                )
                self.interesting_processes[proc.pid] = proc
                self.protect_from_scale_in(protect=True)
                self.dynamodb_atomic_decrement()
            if proc.name() == "Runner.Listener":
                listener_found = True

        if self.in_termating_lifecycle and not listener_found:
            log.info("Runner.Listener process not found - OkayToTerminate instance")
            complete_asg_lifecycle_hook('OkayToTerminate')

    def check_still_alive(self):
        # Check ASG status
        if not self.in_termating_lifecycle:
            state = get_lifecycle_state()
            if state == 'Terminating:Wait':
                self.in_termating_lifecycle = True
                self.gracefully_terminate_runner()
            elif state == 'Pending:Wait':
                complete_asg_lifecycle_hook()

        # proc_connector is un-reliable (UDP) so periodically check if the processes are still alive
        if not self.interesting_processes:
            self.pgrep()
            return

        # list() is used to prevent "Dict changed size during iteration" during loop below
        pids = list(self.interesting_processes.keys())
        log.info("Checking processes %r are still alive", pids)

        for pid in pids:
            proc = self.interesting_processes[pid]
            if not proc.is_running() or proc.status() == psutil.STATUS_ZOMBIE:
                log.info("Proc %d dead but we didn't notice!", pid)
                del self.interesting_processes[pid]

        if not self.interesting_processes:
            log.info("No interesting processes left, unprotecting from scale in")
            self.protect_from_scale_in(protect=False)
        elif not self.protected:
            # If we didn't manage to protect last time, try again
            self.protect_from_scale_in()

    def gracefully_terminate_runner(self):
        check_call(['systemctl', 'stop', 'actions.runner', '--no-block'])

    def protect_from_scale_in(self, protect: bool = True):
        """ Set (or unset) ProtectedFromScaleIn on our instance"""
        if not OWN_ASG:
            # Not part of an ASG
            return

        if self.in_termating_lifecycle:
            log.info("Not trying to SetInstanceProtection, we are already in the terminating lifecycle step")
            return

        asg_client = boto3.client('autoscaling')
        try:
            self._protect_from_scale_in(asg_client, protect)
            self.protected = protect
        except asg_client.exceptions.ClientError as e:
            # This can happen if this the runner picks up a job "too quick", and the ASG still has the state
            # as Pending:Proceed, so we can't yet set it as protected
            log.warning("Failed to set scale in protection: %s", str(e))

    @retry(
        wait=wait_exponential(multiplier=1, max=10),
        stop=stop_after_delay(30),
        before_sleep=before_sleep_log(log, logging.INFO),
        reraise=True,
    )
    def _protect_from_scale_in(self, asg_client, protect):
        asg_client.set_instance_protection(
            AutoScalingGroupName=OWN_ASG,
            InstanceIds=[INSTANCE_ID],
            ProtectedFromScaleIn=protect,
        )

    def dynamodb_atomic_decrement(self):
        dynamodb = boto3.client('dynamodb')
        try:
            resp = dynamodb.update_item(
                TableName=TABLE_NAME,
                Key={'id': {'S': 'queued_jobs'}},
                ExpressionAttributeValues={':delta': {'N': '-1'}, ':limit': {'N': '0'}},
                UpdateExpression='ADD queued :delta',
                # Make sure it never goes below zero!
                ConditionExpression='queued > :limit',
                ReturnValues='UPDATED_NEW',
            )

            log.info("Updated DynamoDB queue length: %s", resp['Attributes']['queued']['N'])
        except dynamodb.exceptions.ConditionalCheckFailedException:
            log.warning("%s.queued was already 0, we won't decrease it any further!", TABLE_NAME)

    def handle_proc_event(self, sock, mask):
        try:
            data, (nlpid, nlgrps) = sock.recvfrom(1024)
        except OSError as e:
            if e.errno == errno.ENOBUFS:
                return
            raise
        if nlpid != 0:
            # Ignore messages from non-root processes
            return

        event, detail = proc_event.from_netlink_packet(data)
        if event.what == ProcEventWhat.EXEC:
            try:
                proc = psutil.Process(detail.pid)

                with proc.oneshot():
                    if proc.name() == "Runner.Worker":
                        log.info(
                            "Found new interesting processes, protecting from scale in %d: %s",
                            detail.pid,
                            proc.cmdline(),
                        )
                        self.interesting_processes[detail.pid] = proc
                        self.protect_from_scale_in(protect=True)
                        self.dynamodb_atomic_decrement()

            except psutil.NoSuchProcess:
                # We lost the race, process has already exited. If it was that short lived it wasn't that
                # interesting anyway
                pass
        elif event.what == ProcEventWhat.EXIT:
            if detail.pid in self.interesting_processes:
                log.info("Interesting process %d exited", detail.pid)
                del self.interesting_processes[detail.pid]

                if not self.interesting_processes:
                    log.info("Watching no processes, disabling termination protection")
                    self.protect_from_scale_in(protect=False)
            elif self.in_termating_lifecycle:
                try:
                    proc = psutil.Process(detail.pid)
                    if proc.name() == "Runner.Listener":
                        log.info("Runner.Listener process exited - OkayToTerminate instance")
                        complete_asg_lifecycle_hook('OkayToTerminate')
                except psutil.NoSuchProcess:
                    # We lost the race, process has already exited. If it was that short lived it wasn't that
                    # interesting anyway
                    pass

    def open_proc_connector_socket(self) -> socket.socket:
        """Open and set up a socket connected to the kernel's Proc Connector event stream

        This uses the Netlink family of socket, the Connector message type and the proc_event connector to get
        send a (UDP) message whenever a process starts or exits.

        We use this mechansim to get notified when processes start or stop, so we can watch for the
        "Runner.Worker" and enable/disable termination protection.
        """

        class bpf_insn(ctypes.Structure):
            """"The BPF instruction data structure"""

            _fields_ = [
                ("code", ctypes.c_ushort),
                ("jt", ctypes.c_ubyte),
                ("jf", ctypes.c_ubyte),
                ("k", ctypes.c_uint32),
            ]

        class bpf_program(ctypes.Structure):
            """"Structure for BIOCSETF"""

            _fields_ = [("bf_len", ctypes.c_uint), ("bf_insns", ctypes.POINTER(bpf_insn))]

            def __init__(self, program):
                self.bf_len = len(program)
                bpf_insn_array = bpf_insn * self.bf_len
                self.bf_insns = bpf_insn_array()

                # Fill the pointer
                for i, insn in enumerate(program):
                    self.bf_insns[i] = insn

        def bpf_jump(code, k, jt, jf) -> bpf_insn:
            """
            :param code: BPF instruction op codes
            :param k: argument
            :param jt: jump offset if true
            :param jf: jump offset if false
            """
            return bpf_insn(code, jt, jf, k)

        def bpf_stmt(code, k):
            return bpf_jump(code, k, 0, 0)

        def packet_filter_prog():
            """
            A Berkley Packet Filter program to filter down the "firehose" of info we receive over the netlink
            socket.

            The Proc Connector doesn't provide any easy way to filter out the firehose of package events, and
            while we could ignore the things we don't care about in Python, it's more efficient to never
            receive those packets. "Luckily" there is the BPF, or Berkley Packet Filter, which can operate on
            any socket. This BPF program was taken from
            https://web.archive.org/web/20130601175512/https://netsplit.com/2011/02/09/the-proc-connector-and-socket-filters/
            """
            # A subset of Berkeley Packet Filter constants and macros, as defined in linux/filter.h.

            # Instruction classes
            BPF_LD = 0x00
            BPF_JMP = 0x05
            BPF_RET = 0x06

            # ld/ldx fields
            BPF_W = 0x00
            BPF_H = 0x08
            BPF_ABS = 0x20

            # alu/jmp fields
            BPF_JEQ = 0x10
            BPF_K = 0x00

            return bpf_program(
                [
                    # Load 16-bit ("half"-word) nlmsg.type field
                    bpf_stmt(BPF_LD | BPF_H | BPF_ABS, NLMsgHdr.type.offset),
                    bpf_jump(BPF_JMP | BPF_JEQ | BPF_K, socket.htons(NlMsgFlag.Done), 1, 0),
                    # Not NlMsgFlag.Done, return whole packet
                    bpf_stmt(BPF_RET | BPF_K, 0xFFFFFFFF),
                    #
                    # Load 32-bit (word) cb_id_idx field
                    bpf_stmt(BPF_LD | BPF_W | BPF_ABS, ctypes.sizeof(NLMsgHdr) + cn_msg.cb_id_idx.offset),
                    bpf_jump(BPF_JMP | BPF_JEQ | BPF_K, socket.htonl(cn_msg.CN_IDX_PROC), 1, 0),
                    # If not CN_IDX_PROC, return whole packet
                    bpf_stmt(BPF_RET | BPF_K, 0xFFFFFFFF),
                    #
                    # Load cb_id_val field
                    bpf_stmt(BPF_LD | BPF_W | BPF_ABS, ctypes.sizeof(NLMsgHdr) + cn_msg.cb_id_val.offset),
                    bpf_jump(BPF_JMP | BPF_JEQ | BPF_K, socket.htonl(cn_msg.CN_VAL_PROC), 1, 0),
                    # If not CN_VAL_PROC, return whole packet
                    bpf_stmt(BPF_RET | BPF_K, 0xFFFFFFFF),
                    #
                    # If not ProcEventWhat.EXEC or ProcEventWhat.EXIT, event, filter out the packet
                    bpf_stmt(
                        BPF_LD | BPF_W | BPF_ABS,
                        ctypes.sizeof(NLMsgHdr) + ctypes.sizeof(cn_msg) + proc_event.what.offset,
                    ),
                    bpf_jump(BPF_JMP | BPF_JEQ | BPF_K, socket.htonl(ProcEventWhat.EXEC), 2, 0),
                    bpf_jump(BPF_JMP | BPF_JEQ | BPF_K, socket.htonl(ProcEventWhat.EXIT), 1, 0),
                    bpf_stmt(BPF_RET | BPF_K, 0x0),
                    # Return everything
                    bpf_stmt(BPF_RET | BPF_K, 0xFFFFFFFF),
                ]
            )

        # Create Netlink socket

        # Missing from most/all pythons
        NETLINK_CONNECTOR = getattr(socket, "NETLINK_CONNECTOR", 11)
        SO_ATTACH_FILTER = getattr(socket, "SO_ATTACH_FILTER", 26)

        sock = socket.socket(socket.AF_NETLINK, socket.SOCK_DGRAM, NETLINK_CONNECTOR)

        filter_prog = packet_filter_prog()
        sock.setsockopt(socket.SOL_SOCKET, SO_ATTACH_FILTER, bytes(filter_prog))  # type: ignore

        sock.bind((os.getpid(), cn_msg.CN_IDX_PROC))

        # Send PROC_CN_MCAST_LISTEN to start receiving messages
        msg = cn_msg(
            header=NLMsgHdr(type=NlMsgFlag.Done, pid=os.getpid()),
            cb_id_idx=cn_msg.CN_IDX_PROC,
            cb_id_val=cn_msg.CN_VAL_PROC,
            seq=0,
            ack=0,
            data=ctypes.c_uint32(ProcConnectorOp.MCAST_LISTEN),
        )

        data = msg.to_bytes()
        if sock.send(data) != len(data):
            raise RuntimeError("Failed to send PROC_CN_MCAST_LISTEN")

        return sock


if __name__ == "__main__":
    main()