xref: /freebsd/tests/atf_python/utils.py (revision 6332ef8941999b0c074d1ece0e1e108447c70b98)
13873bdc2SAlexander V. Chernikov#!/usr/bin/env python3
23873bdc2SAlexander V. Chernikovimport os
3*6332ef89SAlexander V. Chernikovimport pwd
43873bdc2SAlexander V. Chernikovfrom ctypes import CDLL
53873bdc2SAlexander V. Chernikovfrom ctypes import get_errno
63873bdc2SAlexander V. Chernikovfrom ctypes.util import find_library
7*6332ef89SAlexander V. Chernikovfrom typing import Dict
83873bdc2SAlexander V. Chernikovfrom typing import List
93873bdc2SAlexander V. Chernikovfrom typing import Optional
103873bdc2SAlexander V. Chernikov
113873bdc2SAlexander V. Chernikovimport pytest
123873bdc2SAlexander V. Chernikov
133873bdc2SAlexander V. Chernikov
143873bdc2SAlexander V. Chernikovclass LibCWrapper(object):
153873bdc2SAlexander V. Chernikov    def __init__(self):
163873bdc2SAlexander V. Chernikov        path: Optional[str] = find_library("c")
173873bdc2SAlexander V. Chernikov        if path is None:
183873bdc2SAlexander V. Chernikov            raise RuntimeError("libc not found")
193873bdc2SAlexander V. Chernikov        self._libc = CDLL(path, use_errno=True)
203873bdc2SAlexander V. Chernikov
213873bdc2SAlexander V. Chernikov    def modfind(self, mod_name: str) -> int:
223873bdc2SAlexander V. Chernikov        if self._libc.modfind(bytes(mod_name, encoding="ascii")) == -1:
233873bdc2SAlexander V. Chernikov            return get_errno()
243873bdc2SAlexander V. Chernikov        return 0
253873bdc2SAlexander V. Chernikov
263873bdc2SAlexander V. Chernikov    def jail_attach(self, jid: int) -> int:
273873bdc2SAlexander V. Chernikov        if self._libc.jail_attach(jid) != 0:
283873bdc2SAlexander V. Chernikov            return get_errno()
293873bdc2SAlexander V. Chernikov        return 0
303873bdc2SAlexander V. Chernikov
313873bdc2SAlexander V. Chernikov
323873bdc2SAlexander V. Chernikovlibc = LibCWrapper()
333873bdc2SAlexander V. Chernikov
343873bdc2SAlexander V. Chernikov
353873bdc2SAlexander V. Chernikovclass BaseTest(object):
36*6332ef89SAlexander V. Chernikov    NEED_ROOT: bool = False  # True if the class needs root privileges for the setup
37*6332ef89SAlexander V. Chernikov    TARGET_USER = None  # Set to the target user by the framework
383873bdc2SAlexander V. Chernikov    REQUIRED_MODULES: List[str] = []
393873bdc2SAlexander V. Chernikov
403873bdc2SAlexander V. Chernikov    def _check_modules(self):
413873bdc2SAlexander V. Chernikov        for mod_name in self.REQUIRED_MODULES:
423873bdc2SAlexander V. Chernikov            error_code = libc.modfind(mod_name)
433873bdc2SAlexander V. Chernikov            if error_code != 0:
443873bdc2SAlexander V. Chernikov                err_str = os.strerror(error_code)
453873bdc2SAlexander V. Chernikov                pytest.skip(
463873bdc2SAlexander V. Chernikov                    "kernel module '{}' not available: {}".format(mod_name, err_str)
473873bdc2SAlexander V. Chernikov                )
48*6332ef89SAlexander V. Chernikov    @property
49*6332ef89SAlexander V. Chernikov    def atf_vars(self) -> Dict[str, str]:
50*6332ef89SAlexander V. Chernikov        px = "_ATF_VAR_"
51*6332ef89SAlexander V. Chernikov        return {k[len(px):]: v for k, v in os.environ.items() if k.startswith(px)}
52*6332ef89SAlexander V. Chernikov
53*6332ef89SAlexander V. Chernikov    def drop_privileges_user(self, user: str):
54*6332ef89SAlexander V. Chernikov        uid = pwd.getpwnam(user)[2]
55*6332ef89SAlexander V. Chernikov        print("Dropping privs to {}/{}".format(user, uid))
56*6332ef89SAlexander V. Chernikov        os.setuid(uid)
57*6332ef89SAlexander V. Chernikov
58*6332ef89SAlexander V. Chernikov    def drop_privileges(self):
59*6332ef89SAlexander V. Chernikov        if self.TARGET_USER:
60*6332ef89SAlexander V. Chernikov            if self.TARGET_USER == "unprivileged":
61*6332ef89SAlexander V. Chernikov                user = self.atf_vars["unprivileged-user"]
62*6332ef89SAlexander V. Chernikov            else:
63*6332ef89SAlexander V. Chernikov                user = self.TARGET_USER
64*6332ef89SAlexander V. Chernikov            self.drop_privileges_user(user)
653873bdc2SAlexander V. Chernikov
66f63825ffSAlexander V. Chernikov    @property
67*6332ef89SAlexander V. Chernikov    def test_id(self) -> str:
68f63825ffSAlexander V. Chernikov        # 'test_ip6_output.py::TestIP6Output::test_output6_pktinfo[ipandif] (setup)'
69f63825ffSAlexander V. Chernikov        return os.environ.get("PYTEST_CURRENT_TEST").split(" ")[0]
70f63825ffSAlexander V. Chernikov
71f63825ffSAlexander V. Chernikov    def setup_method(self, method):
72f63825ffSAlexander V. Chernikov        """Run all pre-requisits for the test execution"""
733873bdc2SAlexander V. Chernikov        self._check_modules()
74