utils/os_util.py

# Copyright 2023 The ChromiumOS Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

"""Utilities around the os module."""

import functools
import os
from pathlib import Path
import sys
from typing import Optional


def is_root_user() -> bool:
    """Returns True if the user has root privileges.

    For a given process there are two ID's, that we care about. The real user ID
    and effective user ID. The real user ID or simply referred as uid, is the ID
    assigned for the user on whose behalf the process is running. Effective user
    ID is used for privilege checks.

    For a given process the real user ID and effective user ID can be different
    and the access to resources are determined based on the effective user ID.
    For example, a regular user with uid 12345, may not have access to certain
    resources. Running with sudo privileges will make the euid to be 0 (Root)
    (while the uid remains the same 12345) and will gain certain resource
    access.

    Hence to check if a user has root privileges, it is best to check the euid
    of the process.
    """
    return os.geteuid() == 0


def is_non_root_user() -> bool:
    """Returns True if user doesn't have root privileges."""
    return not is_root_user()


def assert_root_user(name: Optional[str] = None):
    """Assert root user."""
    name = name or Path(sys.argv[0]).name
    assert is_root_user(), f"{name}: please run as root user"


def assert_non_root_user(name: Optional[str] = None):
    """Assert root user."""
    name = name or Path(sys.argv[0]).name
    assert is_non_root_user(), f"{name}: please run as non root user"


def require_root_user(_reason):
    """Decorator to note/assert a function must be called as the root user."""

    def outer(func):
        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            assert_root_user(func.__name__)
            return func(*args, **kwargs)

        return wrapper

    return outer


def require_non_root_user(_reason):
    """Decorator to note/assert a function must be called as a non-root user."""

    def outer(func):
        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            assert_non_root_user(func.__name__)
            return func(*args, **kwargs)

        return wrapper

    return outer