tests/atf_python/utils.py

   1 #!/usr/bin/env python3
   2 import os
   3 import pwd
   4 from ctypes import CDLL
   5 from ctypes import get_errno
   6 from ctypes.util import find_library
   7 from typing import Dict
   8 from typing import List
   9 from typing import Optional
  10
  11 import pytest
  12
  13
  14 def 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
  19 class 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
  42 libc = LibCWrapper()
  43
  44
  45 class 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()