#!/usr/bin/env python3
# SPDX-FileCopyrightText: (c) Jonathan Hettwer (bauen1) <j2468h@gmail.com>
# SPDX-License-Identifier: GPL-3.0-only
import sys
if sys.version_info < (3, 10):
    # Due to typing hints and dataclasses kw_only
    sys.exit("Script requires python3.10 or greater")

from typing import Optional, Any, Literal
import re
import argparse
import os
from dataclasses import dataclass, field

include_pattern = re.compile(r'^# ?include "([^"]*)"$')
ioctl_pattern = re.compile(r'{ "([^"]*)", "([^"]*)", ([^,]*), ([^,]*), (.*) },')

# Map of CIL class / common to ioctl prefixes
ioctlname_to_cil_classname: dict[str, tuple[str, ...]] = {
    ## Specific class
    "chr_file": ("RTC_", "DRM_IOCTL_", "USBDEVFS_", "SNDRV_",
        "TCG", "TCS", "TIO", "VT_", # tty/pty ioctls.
        "EVIO",
        "VIDIOC_",
        "VFIO_",
        "SPI_",
        "VHOST_",
        "UI_", # uinput
        "AUTOFS_DEV_IOCTL_",
        "TUN",
        "DM_", # device manager
        "LOOP_",
        "FBIO",
        "WDIOC",
        "AMDKFD_IOC_",
    ),
    "blk_file": (
        "BLK",
        "CDROM",
    ),

    ## Common
    "common_file": (
        "FS_IOC_",
        "BTRFS_",
        "FAT_IOCTL_",
        "VFAT_IOCTL_",
    ),

    # technically by far not all socket ioctls are shared only some are, see net/socket.c:sock_ioctl and the other methods
    "common_socket": ("SIOC", ),
}

cil_common_classes = ["common_file", "common_socket"]

# A map of all CIL classes to their common classes (if any), sorted alphabetically
cil_classes: dict[str, list[Literal["common_file", "common_socket"]]] = {
    "alg_socket": ["common_socket"],
    "appletalk_socket": ["common_socket"],
    "atmpvc_socket": ["common_socket"],
    "ax25_socket": ["common_socket"],
    "blk_file": ["common_file"],
    "bluetooth_socket": ["common_socket"],
    "caif_socket": ["common_socket"],
    "can_socket": ["common_socket"],
    "chr_file": ["common_file"],
    "dccp_socket": ["common_socket"],
    "decnet_socket": ["common_socket"],
    "dir": ["common_file"],
    "fifo_file": ["common_file"],
    "file": ["common_file"],
    "icmp_socket": ["common_socket"],
    "ieee802154_socket": ["common_socket"],
    "ipx_socket": ["common_socket"],
    "irda_socket": ["common_socket"],
    "isdn_socket": ["common_socket"],
    "iucv_socket": ["common_socket"],
    "kcm_socket": ["common_socket"],
    "key_socket": ["common_socket"],
    "llc_socket": ["common_socket"],
    "lnk_file": ["common_file"],
    "mctp_socket": ["common_socket"],
    "netlink_audit_socket": ["common_socket"],
    "netlink_connector_socket": ["common_socket"],
    "netlink_crypto_socket": ["common_socket"],
    "netlink_dnrt_socket": ["common_socket"],
    "netlink_fib_lookup_socket": ["common_socket"],
    "netlink_generic_socket": ["common_socket"],
    "netlink_iscsi_socket": ["common_socket"],
    "netlink_kobject_uevent_socket": ["common_socket"],
    "netlink_netfilter_socket": ["common_socket"],
    "netlink_nflog_socket": ["common_socket"],
    "netlink_rdma_socket": ["common_socket"],
    "netlink_route_socket": ["common_socket"],
    "netlink_scsitransport_socket": ["common_socket"],
    "netlink_selinux_socket": ["common_socket"],
    "netlink_socket": ["common_socket"],
    "netlink_tcpdiag_socket": ["common_socket"],
    "netlink_xfrm_socket": ["common_socket"],
    "netrom_socket": ["common_socket"],
    "nfc_socket": ["common_socket"],
    "packet_socket": ["common_socket"],
    "phonet_socket": ["common_socket"],
    "pppox_socket": ["common_socket"],
    "qipcrtr_socket": ["common_socket"],
    "rawip_socket": ["common_socket"],
    "rds_socket": ["common_socket"],
    "rose_socket": ["common_socket"],
    "rxrpc_socket": ["common_socket"],
    "sctp_socket": ["common_socket"],
    "smc_socket": ["common_socket"],
    "sock_file": ["common_file"],
    "socket": ["common_socket"],
    "tcp_socket": ["common_socket"],
    "tipc_socket": ["common_socket"],
    "tun_socket": ["common_socket"],
    "udp_socket": ["common_socket"],
    "unix_dgram_socket": ["common_socket"],
    "unix_stream_socket": ["common_socket"],
    "vsock_socket": ["common_socket"],
    "x25_socket": ["common_socket"],
    "xdp_socket": ["common_socket"],
}

@dataclass(frozen=True, kw_only=True)
class IOCTLDefinition:
    name: str
    filename: str
    direction: str
    ioctlcmd: int
    size: str
    classname: Optional[str] = field(init=False)

    def __post_init__(self) -> None:
        """After initialisation of self.name we can derive self.classname"""
        for classname, prefixes in ioctlname_to_cil_classname.items():
            if self.name.startswith(prefixes):
                object.__setattr__(self, 'classname', classname) # Bypass frozen=True
                break
        else:
            object.__setattr__(self, 'classname', None) # Bypass frozen=True

    def to_permissionx_statement(self, verbose: int = 0) -> str:
        """Returns the CIL permissionx statement"""
        s = f"(permissionx {self.name} (ioctl {self.classname} (0x{self.ioctlcmd:x})))"
        if verbose >= 1:
            return s + f" ; From {self.filename}"
        else:
            return s

    def to_m4a_define(self, verbose: int = 0) -> str:
        """Returns the refpolicy / m4a macro definition"""
        s = f"define(`{self.name}', `0x{self.ioctlcmd:x}')"
        if verbose >= 1:
            return s + f" # From {self.filename}"
        else:
            return s

def extract_ioctls_from_file(path: str, verbose: int = 0) -> list[IOCTLDefinition]:
    if verbose >= 1:
        print(f"Extracing ioctls from file '{path}'", file=sys.stderr)

    ioctls: list[IOCTLDefinition] = []

    file = open(path, 'r')

    incomment = False
    inpreprocessorcomment = False
    for line in file:
        # Ignore comment lines (fragile)
        if line.startswith('/*'):
            incomment = True

        if line.find("*/") != -1:
            incomment = False
            continue

        if incomment:
            continue
        if line == '\n':
            continue

        # Sort of implement a few preprocessor directives
        if line.startswith(('#include "', '# include "')) and line.endswith('"\n'):
            # We assume a relative path to the file we're looking at:
            matches = include_pattern.match(line)
            if not matches:
                raise ValueError(f"Invalid include in file '{path}': '{line}'")
            include_path = matches[1]
            include_base = os.path.dirname(path)
            actual_path = os.path.join(include_base, include_path)
            actual_path_normed = os.path.normpath(actual_path)
            if verbose >= 1:
                print(f"Following include in '{path}' to '{include_path}' resolved to '{actual_path_normed}'", file=sys.stderr)
            ioctls += extract_ioctls_from_file(actual_path_normed, verbose=verbose)
            continue
        elif line == '#if defined M68K\n':
            inpreprocessorcomment = True
        elif line == '#elif defined X86_64 || defined X32 \\\n':
            continue
        elif line == '#endif\n':
            inpreprocessorcomment = False
            continue

        if inpreprocessorcomment:
            continue

        matches = ioctl_pattern.match(line)
        if matches is None:
            raise ValueError(f'Invalid line in file {path}: {line=}')
        ioctls.append(IOCTLDefinition(name=matches.group(2),
            filename=matches.group(1),
            direction=matches.group(3),
            ioctlcmd=int(matches.group(4), base=16),
            size=matches.group(5)))

    return ioctls

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Generate permissionx statements for ioctls")
    parser.add_argument("type", type=str, choices=["refpolicy", "cil"])
    parser.add_argument("ioctl_files", type=str, nargs='+', help="strace ioctl_inc headers to process")
    parser.add_argument("--verbose", "-v", action="count", default=0, help="Verbosity level")
    parser.add_argument("--cil-namespace", type=str, help="Name of the namespace to use")
    args = parser.parse_args()

    ioctls_parsed = []

    # The files we get are technically C headers, but they will only have /**/ comments and #include
    for file in args.ioctl_files:
        ioctls_parsed += extract_ioctls_from_file(file, verbose=args.verbose)

    ioctls: list[IOCTLDefinition] = []

    # eliminate duplicated definitions
    for ioctl in ioctls_parsed:
        for c in ioctls:
            if ioctl.name == c.name:
                if ioctl.ioctlcmd != c.ioctlcmd:
                    raise ValueError(f"Conflicting ioctl names: {ioctl} and {c} if names match, then ioctlcmd must match!")
                break
        else:
            ioctls.append(ioctl)

    if args.type == "refpolicy":
        # Easy, just dump every ioctl

        # Sort alphabetically first
        ioctls.sort(key=lambda v: v.name)

        for ioctl in ioctls:
            print(ioctl.to_m4a_define(verbose=args.verbose))
    else:
        # Each ioctl definition already resolves to the CIL classname if possible
        # Use this to generate a class -> ioctls mapping

        # TODO: seperate non-mapped and ignored ioctl definitions

        class_to_ioctls: dict[str, list[IOCTLDefinition]] = {}

        for ioctl in ioctls:
            # FIXME: ioctl.classname can be None
            class_to_ioctls.setdefault(ioctl.classname, []).append(ioctl) # type: ignore

        # Sort ioctls by name
        for class_ioctls in class_to_ioctls.values():
            class_ioctls.sort(key=lambda v: v.name)
        # And sort the classes by name
        cil_common_classes.sort()

        # Generate the actual CIL
        print("; Generate by strace_ioctls_inc_to_cil.py")

        indent = ""
        skip_first_newline = False

        if args.cil_namespace:
            print(f"(block {args.cil_namespace}")
            skip_first_newline = True
            indent = "  "

        # First generate a macro for every common class that defines all common ioctls
        # And can be called, e.g. (call generate_common_file_ioctls (sock_file)) to generate e.g.
        # (permissionx TCGETS (ioctl sock_file (0x5401)))

        for common_class in cil_common_classes:
            if skip_first_newline:
                skip_first_newline = False
            else:
                print("")

            print(indent + f"(macro generate_{common_class}_ioctls ((class {common_class}))")
            for ioctl in class_to_ioctls[common_class]:
                print(indent + "  " + ioctl.to_permissionx_statement(verbose=args.verbose))
            print(indent + ")")

        # Now generate a block for every actual class, calling the macros as nescessary
        for classname, commons in cil_classes.items():
            print("")
            print(indent + f"(block {classname}")
            if len(commons) >= 1:
                print(indent + "  ; commons")
                for common in commons:
                    print(indent + f"  (call generate_{common}_ioctls ({classname}))")
            if classname in class_to_ioctls:
                print(indent + "  ; specific ioctls")
                for ioctl in class_to_ioctls[classname]:
                    print(indent + "  " + ioctl.to_permissionx_statement(verbose=args.verbose))
            print(indent + ")")

        if args.cil_namespace:
            print(")")
            indent = ""

        if args.verbose >= 1:
            print("")
            count_missing = len(class_to_ioctls[None]) # type: ignore

            print(f"; Found a total of {len(ioctls)} ioctl definitions with {count_missing} without mapping!")

        if args.verbose >= 2:
            print(f"; The following ioctls are missing a CIL class mapping in the tool:")
            for ioctl in class_to_ioctls[None]: # type: ignore
                print(f"; {ioctl.name} {ioctl.ioctlcmd} defined in file {ioctl.filename}")
