1#!/usr/bin/env python3 2import os 3import pwd 4from ctypes import CDLL 5from ctypes import get_errno 6from ctypes.util import find_library 7from typing import Dict 8from typing import List 9from typing import Optional 10 11import pytest 12 13 14def nodeid_to_method_name(nodeid: str) -> str: 15 """file_name.py::ClassName::method_name[parametrize] -> method_name""" 16 return nodeid.split("::")[-1].split("[")[0] 17 18 19class LibCWrapper(object): 20 def __init__(self): 21 path: Optional[str] = find_library("c") 22 if path is None: 23 raise RuntimeError("libc not found") 24 self._libc = CDLL(path, use_errno=True) 25 26 def modfind(self, mod_name: str) -> int: 27 if self._libc.modfind(bytes(mod_name, encoding="ascii")) == -1: 28 return get_errno() 29 return 0 30 31 def kldload(self, kld_name: str) -> int: 32 if self._libc.kldload(bytes(kld_name, encoding="ascii")) == -1: 33 return get_errno() 34 return 0 35 36 def jail_attach(self, jid: int) -> int: 37 if self._libc.jail_attach(jid) != 0: 38 return get_errno() 39 return 0 40 41 42libc = LibCWrapper() 43 44 45class BaseTest(object): 46 NEED_ROOT: bool = False # True if the class needs root privileges for the setup 47 TARGET_USER = None # Set to the target user by the framework 48 REQUIRED_MODULES: List[str] = [] 49 50 def require_module(self, mod_name: str, skip=True): 51 error_code = libc.modfind(mod_name) 52 if error_code == 0: 53 return 54 err_str = os.strerror(error_code) 55 txt = "kernel module '{}' not available: {}".format(mod_name, err_str) 56 if skip: 57 pytest.skip(txt) 58 else: 59 raise ValueError(txt) 60 61 def _check_modules(self): 62 for mod_name in self.REQUIRED_MODULES: 63 self.require_module(mod_name) 64 65 @property 66 def atf_vars(self) -> Dict[str, str]: 67 px = "_ATF_VAR_" 68 return {k[len(px):]: v for k, v in os.environ.items() if k.startswith(px)} 69 70 def drop_privileges_user(self, user: str): 71 uid = pwd.getpwnam(user)[2] 72 print("Dropping privs to {}/{}".format(user, uid)) 73 os.setuid(uid) 74 75 def drop_privileges(self): 76 if self.TARGET_USER: 77 if self.TARGET_USER == "unprivileged": 78 user = self.atf_vars["unprivileged-user"] 79 else: 80 user = self.TARGET_USER 81 self.drop_privileges_user(user) 82 83 @property 84 def test_id(self) -> str: 85 # 'test_ip6_output.py::TestIP6Output::test_output6_pktinfo[ipandif] (setup)' 86 return os.environ.get("PYTEST_CURRENT_TEST").split(" ")[0] 87 88 def setup_method(self, method): 89 """Run all pre-requisits for the test execution""" 90 self._check_modules() 91