Add test_result_will_be_submitted and get_user_environment function

Author: Nils Bars

ref_utils/__init__.py

@@ -1,8 +1,8 @@
 """Import functions to avoid .dot import for user""" # pylint: disable = invalid-name
-__all__ = ['process', 'assertion', 'checks', 'utils']
+__all__ = ['process', 'assertion', 'utils', "decorator"]
 from .process import drop_privileges, run, get_payload_from_executable, ref_util_install_global_exception_hook, run_with_payload, run_capture_output
 from .assertion import assert_is_dir, assert_is_exec, assert_is_file
-from .utils import print_ok, print_warn, print_err, write_stdout, decode_or_str
+from .utils import print_ok, print_warn, print_err, write_stdout, decode_or_str, test_result_will_be_submitted, get_user_environment
 from .decorator import add_environment_test, add_submission_test, environment_test, submission_test, run_tests, TestResult
 
 ref_util_install_global_exception_hook()

ref_utils/process.py

@@ -14,7 +14,7 @@
 import importlib
 from functools import partial
 
-from .utils import print_err, map_path_as_posix, print_ok, decode_or_str, print_warn
+from .utils import get_user_environment, print_err, map_path_as_posix, print_ok, decode_or_str, print_warn
 from .error import RefUtilsError, RefUtilsProcessTimeoutError, RefUtilsProcessError
 
 _DEFAULT_DROP_UID = 9999
@@ -47,35 +47,6 @@ def ref_util_install_global_exception_hook() -> None:
     hook = partial(ref_util_exception_hook, redact_traceback=False)
     sys.excepthook = hook
 
-
-def get_user_env(last_cmd: Optional[Union[str, bytes]]) -> Dict[str, Union[str, bytes]]:
-    """
-    The task tool (task-wrapper.c) dumps the user environment in the moment it is executed
-    to disk. This function retrives the dumped environment from the file and returns it.
-    This allows to restore the user's exact environment which is paramount for tasks that
-    require a stable stack layout.
-    Returns:
-        The mapping of all key value pairs of the user environment variables that where
-        defined during submission.
-    """
-    ret: Dict[str, Union[str, bytes]] = {}
-    content = Path('/tmp/.user_environ').read_text()
-    lines = content.split('\x00')
-    for line in lines:
-        if line == '':
-            continue
-
-        try:
-            k, v = line.split('=', 1)
-        except Exception as e:
-            print_err(f'Unexpected error while processing "{line}". Error: {e}.')
-        else:
-            ret[k] = v
-
-    if last_cmd is not None:
-        ret['_'] = last_cmd
-    return ret
-
 # Hopefully safe, if not, please tell us, dont mess with the system. Thanks :)
 class RestrictedUnpickler(pickle.Unpickler):
     ALLOWED_MODULE_NAME = {
@@ -169,7 +140,10 @@ def run(cmd_: List[Union[str, Path, bytes]], *args: str, **kwargs: Any) -> 'subp
         # Restore the environment from the user as of the time she called `task ...`.
         # NOTE: The stored environment contains user controlled input!
         # Never restore the environment in a privileged context.
-        kwargs['env'] = get_user_env(cmd[0])
+        env = get_user_environment()
+        # Set the last executed command variable ("_") to the correct value.
+        env["_"] = cmd[0]
+        kwargs['env'] = env
 
     if 'timeout' not in kwargs:
         kwargs['timeout'] = 10
@@ -206,13 +180,15 @@ def run_capture_output(*args: str, check_signal: bool = True, **kwargs: Any) ->
     Wrapper of subprocess.run that redirects stderr to stdout and returns
     (returncode, stdout). This methods raises the same exceptions as ref-utils
     run() method.
+    If `env` is not set, the user environment when she called `task check` is restored.
     """
     p = run(*args, **kwargs, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, check_signal=check_signal)
     return p.returncode, p.stdout
 
 def get_payload_from_executable(cmd_: List[Union[str, Path, bytes]], check: bool = True, check_signal: bool = True, verbose: bool = True, timeout: int = 10) -> Tuple[int, bytes]:
     """
     Get the payload from a script/binary by executing it and returning the output.
+    If `env` is not set, the user environment when she called `task check` is restored.
     Args:
         cmd: The command to execute.
         check: Same as for subprocess.run.

ref_utils/utils.py

@@ -1,5 +1,7 @@
 """ Utility functions including colored printing or subprocess.run wrapper dropping privileges"""
 import sys
+import os
+import typing as t
 from typing import Any, AnyStr, List, Union
 
 from pathlib import Path
@@ -24,18 +26,55 @@ def print_err(*args: str, **kwargs: Any) -> None:
     print(Fore.RED + _sep.join([str(o) for o in args]) + Style.RESET_ALL, **kwargs)
 
 
+def test_result_will_be_submitted() -> bool:
+    """
+    Whether this test execution is going to be submitted as solution.
+    """
+    val = os.environ.get("RESULT_WILL_BE_SUBMITTED")
+    if val:
+        return val.lower() in ["1", "true"]
+    return False
+
+def get_user_environment() -> t.Dict[str, Union[str, bytes]]:
+    """
+    The task tool (task-wrapper.c) dumps the user environment in the moment it is executed
+    to disk. This function retrives the dumped environment from the file and returns it.
+    This allows to restore the user's exact environment which is paramount for tasks that
+    require a stable stack layout.
+    Returns:
+        The mapping of all key value pairs of the user environment variables that where
+        defined during submission.
+    """
+    ret: t.Dict[str, Union[str, bytes]] = {}
+    content = Path('/tmp/.user_environ').read_text()
+    lines = content.split('\x00')
+    for line in lines:
+        if line == '':
+            continue
+
+        try:
+            k, v = line.split('=', 1)
+        except Exception as e:
+            print_err(f'Unexpected error while processing "{line}". Error: {e}.')
+        else:
+            ret[k] = v
+    return ret
+
 def write_stdout(data: AnyStr) -> None:
     sys.stdout.write(data) # type: ignore
 
 
-def decode_or_str(data: bytes) -> str:
+def decode_or_str(data: str | bytes | bytearray) -> str:
     """
     Get a str representing the passed `data`.
     If the bytes are valid UTF8 they are converted to a str.
     Else, they are converted to str via `str(data)`.
     """
     if not data:
         return ''
+    if isinstance(data, str):
+        return data
+
     try:
         return data.decode()
     except: # pylint: disable =bare-except