import ../make-test-python.nix {
  name = "systemd-confinement";

  nodes.machine =
    { pkgs, lib, ... }:
    let
      testLib = pkgs.python3Packages.buildPythonPackage {
        name = "confinement-testlib";
        unpackPhase = ''
          cat > setup.py <<EOF
          from setuptools import setup
          setup(name='confinement-testlib', py_modules=["checkperms"])
          EOF
          cp ${./checkperms.py} checkperms.py
        '';
      };

      mkTest =
        name: testScript:
        pkgs.writers.writePython3 "${name}.py"
          {
            libraries = [
              pkgs.python3Packages.pytest
              testLib
            ];
          }
          ''
            # This runs our test script by using pytest's assertion rewriting, so
            # that whenever we use "assert <something>", the actual values are
            # printed rather than getting a generic AssertionError or the need to
            # pass an explicit assertion error message.
            import ast
            from pathlib import Path
            from _pytest.assertion.rewrite import rewrite_asserts

            script = Path('${pkgs.writeText "${name}-main.py" ''
              import errno, os, pytest, signal
              from subprocess import run
              from checkperms import Accessibility, assert_permissions

              ${testScript}
            ''}') # noqa
            filename = str(script)
            source = script.read_bytes()

            tree = ast.parse(source, filename=filename)
            rewrite_asserts(tree, source, filename)
            exec(compile(tree, filename, 'exec', dont_inherit=True))
          '';

      mkTestStep =
        num:
        {
          description,
          testScript,
          config ? { },
          serviceName ? "test${toString num}",
          rawUnit ? null,
        }:
        {
          systemd.packages = lib.optional (rawUnit != null) (
            pkgs.writeTextFile {
              name = serviceName;
              destination = "/etc/systemd/system/${serviceName}.service";
              text = rawUnit;
            }
          );

          systemd.services.${serviceName} =
            {
              inherit description;
              requiredBy = [ "multi-user.target" ];
              confinement = (config.confinement or { }) // {
                enable = true;
              };
              serviceConfig = (config.serviceConfig or { }) // {
                ExecStart = mkTest serviceName testScript;
                Type = "oneshot";
              };
            }
            // removeAttrs config [
              "confinement"
              "serviceConfig"
            ];
        };

      parametrisedTests =
        lib.concatMap
          (
            { user, privateTmp }:
            let
              withTmp = if privateTmp then "with PrivateTmp" else "without PrivateTmp";

              serviceConfig =
                if user == "static-user" then
                  {
                    User = "chroot-testuser";
                    Group = "chroot-testgroup";
                  }
                else if user == "dynamic-user" then
                  {
                    DynamicUser = true;
                  }
                else
                  { };

            in
            [
              {
                description = "${user}, chroot-only confinement ${withTmp}";
                config = {
                  confinement.mode = "chroot-only";
                  # Only set if privateTmp is true to ensure that the default is false.
                  serviceConfig =
                    serviceConfig
                    // lib.optionalAttrs privateTmp {
                      PrivateTmp = true;
                    };
                };
                testScript =
                  if user == "root" then
                    ''
                      assert os.getuid() == 0
                      assert os.getgid() == 0

                      assert_permissions({
                        'bin': Accessibility.READABLE,
                        'nix': Accessibility.READABLE,
                        'run': Accessibility.READABLE,
                        ${lib.optionalString privateTmp "'tmp': Accessibility.STICKY,"}
                        ${lib.optionalString privateTmp "'var': Accessibility.READABLE,"}
                        ${lib.optionalString privateTmp "'var/tmp': Accessibility.STICKY,"}
                      })
                    ''
                  else
                    ''
                      assert os.getuid() != 0
                      assert os.getgid() != 0

                      assert_permissions({
                        'bin': Accessibility.READABLE,
                        'nix': Accessibility.READABLE,
                        'run': Accessibility.READABLE,
                        ${lib.optionalString privateTmp "'tmp': Accessibility.STICKY,"}
                        ${lib.optionalString privateTmp "'var': Accessibility.READABLE,"}
                        ${lib.optionalString privateTmp "'var/tmp': Accessibility.STICKY,"}
                      })
                    '';
              }
              {
                description = "${user}, full APIVFS confinement ${withTmp}";
                config = {
                  # Only set if privateTmp is false to ensure that the default is true.
                  serviceConfig =
                    serviceConfig
                    // lib.optionalAttrs (!privateTmp) {
                      PrivateTmp = false;
                    };
                };
                testScript =
                  if user == "root" then
                    ''
                      assert os.getuid() == 0
                      assert os.getgid() == 0

                      assert_permissions({
                        'bin': Accessibility.READABLE,
                        'nix': Accessibility.READABLE,
                        ${lib.optionalString privateTmp "'tmp': Accessibility.STICKY,"}
                        'run': Accessibility.WRITABLE,

                        'proc': Accessibility.SPECIAL,
                        'sys': Accessibility.SPECIAL,
                        'dev': Accessibility.WRITABLE,

                        ${lib.optionalString privateTmp "'var': Accessibility.READABLE,"}
                        ${lib.optionalString privateTmp "'var/tmp': Accessibility.STICKY,"}
                      })
                    ''
                  else
                    ''
                      assert os.getuid() != 0
                      assert os.getgid() != 0

                      assert_permissions({
                        'bin': Accessibility.READABLE,
                        'nix': Accessibility.READABLE,
                        ${lib.optionalString privateTmp "'tmp': Accessibility.STICKY,"}
                        'run': Accessibility.STICKY,

                        'proc': Accessibility.SPECIAL,
                        'sys': Accessibility.SPECIAL,
                        'dev': Accessibility.SPECIAL,
                        'dev/shm': Accessibility.STICKY,
                        'dev/mqueue': Accessibility.STICKY,

                        ${lib.optionalString privateTmp "'var': Accessibility.READABLE,"}
                        ${lib.optionalString privateTmp "'var/tmp': Accessibility.STICKY,"}
                      })
                    '';
              }
            ]
          )
          (
            lib.cartesianProduct {
              user = [
                "root"
                "dynamic-user"
                "static-user"
              ];
              privateTmp = [
                true
                false
              ];
            }
          );

    in
    {
      imports = lib.imap1 mkTestStep (
        parametrisedTests
        ++ [
          {
            description = "existence of bind-mounted /etc";
            config.serviceConfig.BindReadOnlyPaths = [ "/etc" ];
            testScript = ''
              assert Path('/etc/passwd').read_text()
            '';
          }
          (
            let
              symlink = pkgs.runCommand "symlink" {
                target = pkgs.writeText "symlink-target" "got me";
              } "ln -s \"$target\" \"$out\"";
            in
            {
              description = "check if symlinks are properly bind-mounted";
              config.confinement.packages = lib.singleton symlink;
              testScript = ''
                assert Path('${symlink}').read_text() == 'got me'
              '';
            }
          )
          {
            description = "check if StateDirectory works";
            config.serviceConfig.User = "chroot-testuser";
            config.serviceConfig.Group = "chroot-testgroup";
            config.serviceConfig.StateDirectory = "testme";

            # We restart on purpose here since we want to check whether the state
            # directory actually persists.
            config.serviceConfig.Restart = "on-failure";
            config.serviceConfig.RestartMode = "direct";

            testScript = ''
              assert not Path('/tmp/canary').exists()
              Path('/tmp/canary').touch()

              if (foo := Path('/var/lib/testme/foo')).exists():
                assert Path('/var/lib/testme/foo').read_text() == 'works'
              else:
                Path('/var/lib/testme/foo').write_text('works')
                print('<4>Exiting with failure to check persistence on restart.')
                raise SystemExit(1)
            '';
          }
          {
            description = "check if /bin/sh works";
            testScript = ''
              assert Path('/bin/sh').exists()

              result = run(
                ['/bin/sh', '-c', 'echo -n bar'],
                capture_output=True,
                check=True,
              )
              assert result.stdout == b'bar'
            '';
          }
          {
            description = "check if suppressing /bin/sh works";
            config.confinement.binSh = null;
            testScript = ''
              assert not Path('/bin/sh').exists()
              with pytest.raises(FileNotFoundError):
                run(['/bin/sh', '-c', 'echo foo'])
            '';
          }
          {
            description = "check if we can set /bin/sh to something different";
            config.confinement.binSh = "${pkgs.hello}/bin/hello";
            testScript = ''
              assert Path('/bin/sh').exists()
              result = run(
                ['/bin/sh', '-g', 'foo'],
                capture_output=True,
                check=True,
              )
              assert result.stdout == b'foo\n'
            '';
          }
          {
            description = "check if only Exec* dependencies are included";
            config.environment.FOOBAR = pkgs.writeText "foobar" "eek";
            testScript = ''
              with pytest.raises(FileNotFoundError):
                Path(os.environ['FOOBAR']).read_text()
            '';
          }
          {
            description = "check if fullUnit includes all dependencies";
            config.environment.FOOBAR = pkgs.writeText "foobar" "eek";
            config.confinement.fullUnit = true;
            testScript = ''
              assert Path(os.environ['FOOBAR']).read_text() == 'eek'
            '';
          }
          {
            description = "check if shipped unit file still works";
            config.confinement.mode = "chroot-only";
            rawUnit = ''
              [Service]
              SystemCallFilter=~kill
              SystemCallErrorNumber=ELOOP
            '';
            testScript = ''
              with pytest.raises(OSError) as excinfo:
                os.kill(os.getpid(), signal.SIGKILL)
              assert excinfo.value.errno == errno.ELOOP
            '';
          }
        ]
      );

      config.users.groups.chroot-testgroup = { };
      config.users.users.chroot-testuser = {
        isSystemUser = true;
        description = "Chroot Test User";
        group = "chroot-testgroup";
      };
    };

  testScript = ''
    machine.wait_for_unit("multi-user.target")
  '';
}