13873bdc2SAlexander V. Chernikov#!/usr/bin/env python3 23873bdc2SAlexander V. Chernikovimport os 36332ef89SAlexander V. Chernikovimport pwd 43873bdc2SAlexander V. Chernikovfrom ctypes import CDLL 53873bdc2SAlexander V. Chernikovfrom ctypes import get_errno 63873bdc2SAlexander V. Chernikovfrom ctypes.util import find_library 76332ef89SAlexander 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 14d9af4219SAlexander V. Chernikovdef nodeid_to_method_name(nodeid: str) -> str: 15d9af4219SAlexander V. Chernikov """file_name.py::ClassName::method_name[parametrize] -> method_name""" 16d9af4219SAlexander V. Chernikov return nodeid.split("::")[-1].split("[")[0] 17d9af4219SAlexander V. Chernikov 18d9af4219SAlexander V. Chernikov 193873bdc2SAlexander V. Chernikovclass LibCWrapper(object): 203873bdc2SAlexander V. Chernikov def __init__(self): 213873bdc2SAlexander V. Chernikov path: Optional[str] = find_library("c") 223873bdc2SAlexander V. Chernikov if path is None: 233873bdc2SAlexander V. Chernikov raise RuntimeError("libc not found") 243873bdc2SAlexander V. Chernikov self._libc = CDLL(path, use_errno=True) 253873bdc2SAlexander V. Chernikov 263873bdc2SAlexander V. Chernikov def modfind(self, mod_name: str) -> int: 273873bdc2SAlexander V. Chernikov if self._libc.modfind(bytes(mod_name, encoding="ascii")) == -1: 283873bdc2SAlexander V. Chernikov return get_errno() 293873bdc2SAlexander V. Chernikov return 0 303873bdc2SAlexander V. Chernikov 313e5d0784SAlexander V. Chernikov def kldload(self, kld_name: str) -> int: 323e5d0784SAlexander V. Chernikov if self._libc.kldload(bytes(kld_name, encoding="ascii")) == -1: 333e5d0784SAlexander V. Chernikov return get_errno() 343e5d0784SAlexander V. Chernikov return 0 353e5d0784SAlexander V. Chernikov 363873bdc2SAlexander V. Chernikov def jail_attach(self, jid: int) -> int: 373873bdc2SAlexander V. Chernikov if self._libc.jail_attach(jid) != 0: 383873bdc2SAlexander V. Chernikov return get_errno() 393873bdc2SAlexander V. Chernikov return 0 403873bdc2SAlexander V. Chernikov 413873bdc2SAlexander V. Chernikov 423873bdc2SAlexander V. Chernikovlibc = LibCWrapper() 433873bdc2SAlexander V. Chernikov 443873bdc2SAlexander V. Chernikov 453873bdc2SAlexander V. Chernikovclass BaseTest(object): 466332ef89SAlexander V. Chernikov NEED_ROOT: bool = False # True if the class needs root privileges for the setup 476332ef89SAlexander V. Chernikov TARGET_USER = None # Set to the target user by the framework 483873bdc2SAlexander V. Chernikov REQUIRED_MODULES: List[str] = [] 493873bdc2SAlexander V. Chernikov 5097760572SAlexander V. Chernikov def require_module(self, mod_name: str, skip=True): 5197760572SAlexander V. Chernikov error_code = libc.modfind(mod_name) 5297760572SAlexander V. Chernikov if error_code == 0: 5397760572SAlexander V. Chernikov return 5497760572SAlexander V. Chernikov err_str = os.strerror(error_code) 5597760572SAlexander V. Chernikov txt = "kernel module '{}' not available: {}".format(mod_name, err_str) 5697760572SAlexander V. Chernikov if skip: 5797760572SAlexander V. Chernikov pytest.skip(txt) 5897760572SAlexander V. Chernikov else: 5997760572SAlexander V. Chernikov raise ValueError(txt) 6097760572SAlexander V. Chernikov 613873bdc2SAlexander V. Chernikov def _check_modules(self): 623873bdc2SAlexander V. Chernikov for mod_name in self.REQUIRED_MODULES: 6397760572SAlexander V. Chernikov self.require_module(mod_name) 6497760572SAlexander V. Chernikov 656332ef89SAlexander V. Chernikov @property 666332ef89SAlexander V. Chernikov def atf_vars(self) -> Dict[str, str]: 676332ef89SAlexander V. Chernikov px = "_ATF_VAR_" 686332ef89SAlexander V. Chernikov return {k[len(px):]: v for k, v in os.environ.items() if k.startswith(px)} 696332ef89SAlexander V. Chernikov 706332ef89SAlexander V. Chernikov def drop_privileges_user(self, user: str): 716332ef89SAlexander V. Chernikov uid = pwd.getpwnam(user)[2] 726332ef89SAlexander V. Chernikov print("Dropping privs to {}/{}".format(user, uid)) 736332ef89SAlexander V. Chernikov os.setuid(uid) 746332ef89SAlexander V. Chernikov 756332ef89SAlexander V. Chernikov def drop_privileges(self): 766332ef89SAlexander V. Chernikov if self.TARGET_USER: 776332ef89SAlexander V. Chernikov if self.TARGET_USER == "unprivileged": 786332ef89SAlexander V. Chernikov user = self.atf_vars["unprivileged-user"] 796332ef89SAlexander V. Chernikov else: 806332ef89SAlexander V. Chernikov user = self.TARGET_USER 816332ef89SAlexander V. Chernikov self.drop_privileges_user(user) 823873bdc2SAlexander V. Chernikov 83f63825ffSAlexander V. Chernikov @property 846332ef89SAlexander V. Chernikov def test_id(self) -> str: 85f63825ffSAlexander V. Chernikov # 'test_ip6_output.py::TestIP6Output::test_output6_pktinfo[ipandif] (setup)' 86f63825ffSAlexander V. Chernikov return os.environ.get("PYTEST_CURRENT_TEST").split(" ")[0] 87f63825ffSAlexander V. Chernikov 88f63825ffSAlexander V. Chernikov def setup_method(self, method): 89f63825ffSAlexander V. Chernikov """Run all pre-requisits for the test execution""" 903873bdc2SAlexander V. Chernikov self._check_modules() 91