#!/usr/bin/env python3
"""Root/capability checks, privilege dropping, directory permission enforcement."""

import ctypes
import ctypes.util
import grp
import os
import pwd
import stat
from pathlib import Path
from typing import Optional

# Linux capability constants
CAP_NET_RAW = 13
CAP_NET_ADMIN = 15
CAP_SYS_ADMIN = 21
CAP_DAC_OVERRIDE = 1

# capget/capset structures
_LINUX_CAPABILITY_VERSION_3 = 0x20080522
_VFS_CAP_REVISION_2 = 0x02000000

_libc = ctypes.CDLL(ctypes.util.find_library("c"), use_errno=True)


class _cap_header(ctypes.Structure):
    _fields_ = [
        ("version", ctypes.c_uint32),
        ("pid", ctypes.c_int),
    ]


class _cap_data(ctypes.Structure):
    _fields_ = [
        ("effective", ctypes.c_uint32),
        ("permitted", ctypes.c_uint32),
        ("inheritable", ctypes.c_uint32),
    ]


def check_root() -> bool:
    return os.geteuid() == 0


def check_capability(cap: int) -> bool:
    if check_root():
        return True

    try:
        header = _cap_header(version=_LINUX_CAPABILITY_VERSION_3, pid=0)
        # Version 3 uses 2 data structs for 64-bit capability sets
        data = (_cap_data * 2)()
        ret = _libc.capget(ctypes.byref(header), data)
        if ret != 0:
            return False

        idx = cap // 32
        bit = 1 << (cap % 32)
        return bool(data[idx].effective & bit)
    except Exception:
        return False


def has_net_raw() -> bool:
    return check_capability(CAP_NET_RAW)


def has_net_admin() -> bool:
    return check_capability(CAP_NET_ADMIN)


def check_required_capabilities() -> dict:
    return {
        "root": check_root(),
        "CAP_NET_RAW": check_capability(CAP_NET_RAW),
        "CAP_NET_ADMIN": check_capability(CAP_NET_ADMIN),
        "CAP_SYS_ADMIN": check_capability(CAP_SYS_ADMIN),
    }


def drop_privileges(
    username: str = "nobody",
    groupname: Optional[str] = None,
    keep_caps: Optional[list] = None,
) -> bool:
    if not check_root():
        return False

    try:
        pw = pwd.getpwnam(username)
        uid = pw.pw_uid
        gid = pw.pw_gid

        if groupname:
            gr = grp.getgrnam(groupname)
            gid = gr.gr_gid

        # Set supplementary groups
        os.setgroups([])

        # Set GID first (can't change after dropping root)
        os.setregid(gid, gid)
        os.setreuid(uid, uid)

        # Verify drop
        if os.geteuid() == 0:
            return False

        # Re-apply specific capabilities if requested
        if keep_caps:
            _set_keepcaps()
            _apply_caps(keep_caps)

        return True
    except (KeyError, PermissionError, OSError):
        return False


def _set_keepcaps():
    PR_SET_KEEPCAPS = 8
    _libc.prctl(PR_SET_KEEPCAPS, 1, 0, 0, 0)


def _apply_caps(caps: list):
    header = _cap_header(version=_LINUX_CAPABILITY_VERSION_3, pid=0)
    data = (_cap_data * 2)()

    for cap in caps:
        idx = cap // 32
        bit = 1 << (cap % 32)
        data[idx].effective |= bit
        data[idx].permitted |= bit

    _libc.capset(ctypes.byref(header), data)


def enforce_directory_permissions(
    path: str,
    mode: int = 0o700,
    recursive: bool = True,
) -> int:
    count = 0
    target = Path(path)

    if not target.exists():
        target.mkdir(parents=True, mode=mode)
        count += 1
        return count

    current = target.stat().st_mode & 0o777
    if current != mode:
        target.chmod(mode)
        count += 1

    if recursive:
        for item in target.rglob("*"):
            if item.is_dir():
                item_mode = item.stat().st_mode & 0o777
                if item_mode != mode:
                    item.chmod(mode)
                    count += 1
            elif item.is_file():
                file_mode = mode & 0o700  # Files get owner-only based on dir mode
                item_mode = item.stat().st_mode & 0o777
                if item_mode != file_mode:
                    item.chmod(file_mode)
                    count += 1

    return count


def secure_path(path: str, owner_uid: Optional[int] = None) -> bool:
    try:
        p = Path(path)
        if not p.exists():
            return False

        # Set ownership if specified
        if owner_uid is not None:
            os.chown(path, owner_uid, owner_uid)

        # Restrict permissions
        if p.is_dir():
            p.chmod(0o700)
        else:
            p.chmod(0o600)

        return True
    except (OSError, PermissionError):
        return False