Capability experiments

caps.py

import ctypes
import ctypes.util
import enum
import errno

from collections.abc import Generator
from collections import OrderedDict
from typing import Dict, Iterable, Optional

import osbuild
import osbuild.meta


class LibCap:

    cap_value_t = ctypes.c_int

    def __init__(self, lib) -> None:
        self.lib = lib

        # process-wide bounding set
        get_bound = lib.cap_get_bound
        get_bound.argtypes = (self.cap_value_t,)
        get_bound.restype = ctypes.c_int
        get_bound.errcheck = self._check_result

        self._get_bound = get_bound

        from_name = lib.cap_from_name
        from_name.argtypes = (ctypes.c_char_p, ctypes.POINTER(self.cap_value_t),)
        from_name.restype = ctypes.c_int
        from_name.errcheck = self._check_result

        self._from_name = from_name

    @staticmethod
    def _check_result(result, func, args):
        if result == -1:
            raise OSError(errno.EINVAL, func, args)
        return result

    @staticmethod
    def make():
        path = ctypes.util.find_library("cap")
        if not path:
            return None

        try:
            lib = ctypes.CDLL(path)
        except (OSError, ImportError):
            return None

        return LibCap(lib)

    def supports(self, capability: str) -> bool:
        try:
            cap = self.cap_value_t()
            self._from_name(capability.encode("utf-8"), ctypes.pointer(cap))
            self._get_bound(cap)
            return True
        except OSError:
            return False


class CapSet:

    def __init__(self,
                 caps: Iterable[str] = None,
                 drop: Iterable[str] = None) -> None:
        self._caps: OrderedDict[str, bool] = OrderedDict(
            zip(caps or [], [True] * len(caps))
        )
        self._caps.update(zip(drop or [], [False] * len(drop)))

    def set(self, capability: str, enabled: bool = True) -> None:
        self._caps[capability] = enabled
        self._caps.move_to_end(capability, last=not enabled)

    def get(self, capability: str) -> Optional[bool]:
        self._caps.get(capability)

    def enable(self, *capabilities: Iterable[str]) -> None:
        _ = [self.set(cap, True) for cap in capabilities]

    def disable(self, *capabilities: Iterable[str]) -> None:
        _ = [self.set(cap, False) for cap in capabilities]

    def discard(self, capability: str) -> None:
        self._caps.pop(capability, None)

    def __len__(self) -> int:
        return len(self._caps)

    def __contains__(self, capability: str) -> bool:
        return capability in self._caps

    def __getitem__(self, item) -> bool:
        return self._caps[item]

    def __iter__(self) -> Generator[str, bool]:
        for cap, val in self._caps.items():
            yield cap, val

    def __str__(self):
        l = [
            f"{'+' if on else '-'}{cap}"
            for cap, on in self
        ]

        return " ".join(l)


if __name__ == "__main__":
    cs = CapSet(caps={"CAP_SETUID", "CAP_SETGID"}, drop={"CAP_MAC_ADMIN"})
    print(str(cs))
    print(repr(cs))
    cs.set("CAP_SETUID", False)
    cs.enable("CAP_PERFMON", "CAP_SYS_PTRACE")
    print(str(cs))

    index = osbuild.meta.Index(".")
    info = index.get_module_info("Stage", "org.osbuild.copy")
    print(info.caps)

    lib = LibCap.make()
    if lib:
        res = lib.supports("CAP_SETUID")
        print("systems supports CAP_SETUID:", res)

        res = lib.supports("CAP_SETUID_BLA")
        print("systems supports CAP_SETUID_BLA:", res)