nixos/taskserver: Rewrite helper-tool in Python

In the comments of the pull request @nbp wrote:

"Why is it implemented in 3 different languages: Nix, Bash and C?"

And he's right, it doesn't make sense, because we were using C as a
runuser replacement and used Nix to generate the shellscript
boilerplates.

Writing this in Python gets rid of all of this and we also don't need
the boilerplate as well, because we're using Click to handle all the
command line stuff.

Note that this currently is a 1:1 implementation of what we had before.

Signed-off-by: aszlig <aszlig@redmoonstudios.org>

3 files changed, 299 insertions, 327 deletions

diff --git a/nixos/modules/services/misc/taskserver/default.nix b/nixos/modules/services/misc/taskserver/default.nix
index 62e35803117c..86eabb9bcfc8 100644
--- a/nixos/modules/services/misc/taskserver/default.nix
+++ b/nixos/modules/services/misc/taskserver/default.nix
@@ -80,9 +80,29 @@ let
 
   mkShellStr = val: "'${replaceStrings ["'"] ["'\\''"] val}'";
 
-  nixos-taskserver = import ./helper-tool.nix {
-    inherit pkgs lib mkShellStr taskd;
-    config = cfg;
+  nixos-taskserver = pkgs.buildPythonPackage {
+    name = "nixos-taskserver";
+    namePrefix = "";
+
+    src = pkgs.runCommand "nixos-taskserver-src" {} ''
+      mkdir -p "$out"
+      cat "${pkgs.substituteAll {
+        src = ./helper-tool.py;
+        certtool = "${pkgs.gnutls}/bin/certtool";
+        inherit taskd;
+        inherit (cfg) dataDir user group;
+        inherit (cfg.server) fqdn;
+      }}" > "$out/main.py"
+      cat > "$out/setup.py" <<EOF
+      from setuptools import setup
+      setup(name="nixos-taskserver",
+            py_modules=["main"],
+            install_requires=["Click"],
+            entry_points="[console_scripts]\\nnixos-taskserver=main:cli")
+      EOF
+    '';
+
+    propagatedBuildInputs = [ pkgs.pythonPackages.click ];
   };
 
   ctlcmd = "${nixos-taskserver}/bin/nixos-taskserver --service-helper";
diff --git a/nixos/modules/services/misc/taskserver/helper-tool.nix b/nixos/modules/services/misc/taskserver/helper-tool.nix
deleted file mode 100644
index 70660574d04c..000000000000
--- a/nixos/modules/services/misc/taskserver/helper-tool.nix
+++ /dev/null
@@ -1,324 +0,0 @@
-{ config, pkgs, lib, mkShellStr, taskd }:
-
-let
-  commandName = "nixos-taskserver";
-  mkShellName = lib.replaceStrings ["-"] ["_"];
-
-  genClientKey = ''
-    umask 0077
-    if tmpdir="$(${pkgs.coreutils}/bin/mktemp -d)"; then
-      trap "rm -rf '$tmpdir'" EXIT
-      ${pkgs.gnutls}/bin/certtool -p --bits 2048 --outfile "$tmpdir/key"
-
-      cat > "$tmpdir/template" <<-\ \ EOF
-      organization = $organisation
-      cn = ${config.server.fqdn}
-      tls_www_client
-      encryption_key
-      signing_key
-      EOF
-
-      ${pkgs.gnutls}/bin/certtool -c \
-        --load-privkey "$tmpdir/key" \
-        --load-ca-privkey "${config.dataDir}/keys/ca.key" \
-        --load-ca-certificate "${config.dataDir}/keys/ca.cert" \
-        --template "$tmpdir/template" \
-        --outfile "$tmpdir/cert"
-
-      mkdir -m 0700 -p "${config.dataDir}/keys/user/$organisation/$user"
-      chown root:root "${config.dataDir}/keys/user/$organisation/$user"
-      cat "$tmpdir/key" \
-        > "${config.dataDir}/keys/user/$organisation/$user/private.key"
-      cat "$tmpdir/cert" \
-        > "${config.dataDir}/keys/user/$organisation/$user/public.cert"
-
-      rm -rf "$tmpdir"
-      trap - EXIT
-    else
-      echo "Unable to create temporary directory for client" \
-           "certificate creation." >&2
-      exit 1
-    fi
-  '';
-
-  mkSubCommand = name: { args, description, script }: let
-    mkArg = pos: arg: "local ${arg}=\"\$${toString pos}\"";
-    mkDesc = line: "echo ${mkShellStr "    ${line}"} >&2";
-    usagePosArgs = lib.concatMapStringsSep " " (a: "<${a}>") args;
-  in ''
-    subcmd_${mkShellName name}() {
-      ${lib.concatImapStringsSep "\n  " mkArg args}
-      ${script}
-    }
-
-    usage_${mkShellName name}() {
-      echo "  ${commandName} ${name} ${usagePosArgs}" >&2
-      ${lib.concatMapStringsSep "\n  " mkDesc description}
-    }
-  '';
-
-  mkCStr = val: "\"${lib.escape ["\\" "\""] val}\"";
-
-  taskdUser = let
-    runUser = pkgs.writeText "runuser.c" ''
-      #include <sys/types.h>
-      #include <pwd.h>
-      #include <grp.h>
-      #include <stdio.h>
-      #include <stdlib.h>
-      #include <errno.h>
-      #include <unistd.h>
-
-      int main(int argc, char **argv) {
-        struct passwd *userinfo;
-        struct group *groupinfo;
-        errno = 0;
-        if ((userinfo = getpwnam(${mkCStr config.user})) == NULL) {
-          if (errno == 0)
-            fputs(${mkCStr "User name `${config.user}' not found."}, stderr);
-          else
-            perror("getpwnam");
-          return EXIT_FAILURE;
-        }
-        errno = 0;
-        if ((groupinfo = getgrnam(${mkCStr config.group})) == NULL) {
-          if (errno == 0)
-            fputs(${mkCStr "Group name `${config.group}' not found."}, stderr);
-          else
-            perror("getgrnam");
-          return EXIT_FAILURE;
-        }
-        if (setgid(groupinfo->gr_gid) == -1) {
-          perror("setgid");
-          return EXIT_FAILURE;
-        }
-        if (setuid(userinfo->pw_uid) == -1) {
-          perror("setgid");
-          return EXIT_FAILURE;
-        }
-        argv[0] = "taskd";
-        if (execv(${mkCStr taskd}, argv) == -1) {
-          perror("execv");
-          return EXIT_FAILURE;
-        }
-        /* never reached */
-        return EXIT_SUCCESS;
-      }
-    '';
-  in pkgs.runCommand "taskd-user" {} ''
-    cc -Wall -std=c11 "${runUser}" -o "$out"
-  '';
-
-  subcommands = {
-    list-users = {
-      args = [ "organisation" ];
-
-      description = [
-        "List all users belonging to the specified organisation."
-      ];
-
-      script = ''
-        legend "The following users exist for $organisation:"
-        ${pkgs.findutils}/bin/find \
-          "${config.dataDir}/orgs/$organisation/users" \
-          -mindepth 2 -maxdepth 2 -name config \
-          -exec ${pkgs.gnused}/bin/sed -ne 's/^user *= *//p' {} +
-      '';
-    };
-
-    list-orgs = {
-      args = [];
-
-      description = [
-        "List available organisations"
-      ];
-
-      script = ''
-        legend "The following organisations exist:"
-        ${pkgs.findutils}/bin/find \
-          "${config.dataDir}/orgs" -mindepth 1 -maxdepth 1 \
-          -type d
-      '';
-    };
-
-    get-uuid = {
-      args = [ "organisation" "user" ];
-
-      description = [
-        "Get the UUID of the specified user belonging to the specified"
-        "organisation."
-      ];
-
-      script = ''
-        for uuid in "${config.dataDir}/orgs/$organisation/users"/*; do
-          usr="$(${pkgs.gnused}/bin/sed -ne 's/^user *= *//p' "$uuid/config")"
-          if [ "$usr" = "$user" ]; then
-            legend "User $user has the following UUID:"
-            echo "$(${pkgs.coreutils}/bin/basename "$uuid")"
-            exit 0
-          fi
-        done
-        echo "No UUID found for user $user." >&2
-        exit 1
-      '';
-    };
-
-    export-user = {
-      args = [ "organisation" "user" ];
-
-      description = [
-        "Export user of the specified organisation as a series of shell"
-        "commands that can be used on the client side to easily import"
-        "the certificates."
-        ""
-        "Note that the private key will be exported as well, so use this"
-        "with care!"
-      ];
-
-      script = ''
-        if ! subcmd_quiet list-users "$organisation" | grep -qxF "$user"; then
-          exists "User $user doesn't exist in organisation $organisation."
-        fi
-
-        uuid="$(subcmd_quiet get-uuid "$organisation" "$user")" || exit 1
-
-        cat <<COMMANDS
-        taskdatadir="\''${TASKDATA:-\$HOME/.task}"
-        umask 0077
-        mkdir -p "\$taskdatadir/keys"
-        cat > "\$taskdatadir/keys/public.cert" <<EOF
-        $(cat "${config.dataDir}/keys/user/$organisation/$user/public.cert")
-        EOF
-        cat > "\$taskdatadir/keys/private.key" <<EOF
-        $(${pkgs.gnused}/bin/sed -ne '/^---* *BEGIN /,/^---* *END /p' \
-          "${config.dataDir}/keys/user/$organisation/$user/private.key")
-        EOF
-        cat > "\$taskdatadir/keys/ca.cert" <<EOF
-        $(cat "${config.dataDir}/keys/ca.cert")
-        EOF
-        task config taskd.certificate -- "\$taskdatadir/keys/public.cert"
-        task config taskd.key         -- "\$taskdatadir/keys/private.key"
-        task config taskd.ca          -- "\$taskdatadir/keys/ca.cert"
-        task config taskd.credentials -- "$organisation/$user/$uuid"
-        COMMANDS
-      '';
-    };
-
-    add-org = {
-      args = [ "name" ];
-
-      description = [
-        "Create an organisation with the specified name."
-      ];
-
-      script = ''
-        if [ -e "orgs/$name" ]; then
-          exists "Organisation with name $name already exists."
-        fi
-        ${taskdUser} add org "$name"
-      '';
-    };
-
-    add-user = {
-      args = [ "organisation" "user" ];
-
-      description = [
-        "Create a user for the given organisation and print the UUID along"
-        "with the client certificate and key."
-      ];
-
-      script = ''
-        if subcmd list-users "$organisation" | grep -qxF "$user"; then
-          exists "User $user already exists in organisation $organisation."
-        fi
-        ${taskdUser} add user "$organisation" "$user"
-        ${genClientKey}
-      '';
-    };
-
-    add-group = {
-      args = [ "organisation" "group" ];
-
-      description = [
-        "Create a group for the given organisation."
-      ];
-
-      script = ''
-        if [ -e "orgs/$organisation/groups/$group" ]; then
-          exists "Group $group already exists in organisation $organisation."
-        fi
-        ${taskdUser} add group "$organisation" "$group"
-      '';
-    };
-  };
-
-  mkCase = name: { args, ... }: let
-    mkPosArg = pos: lib.const "\"\$${toString (pos + 1)}\"";
-    cmdArgs = lib.concatImapStringsSep " " mkPosArg args;
-  in ''
-    ${name})
-      if [ $# -ne ${toString ((lib.length args) + 1)} ]; then
-        echo "Wrong number of arguments to ${name}." >&2
-        echo >&2
-        usage_${mkShellName name}
-        exit 1
-      fi
-      subcmd "${name}" ${cmdArgs};;
-  '';
-
-in pkgs.writeScriptBin commandName ''
-  #!${pkgs.stdenv.shell}
-  export TASKDDATA=${mkShellStr config.dataDir}
-
-  quiet=0
-  # Deliberately undocumented, because we don't want people to use this as
-  # it's only used in and specific to the preStart script of the Taskserver
-  # service.
-  if [ "$1" = "--service-helper" ]; then
-    quiet=1
-    exists() {
-      exit 0
-    }
-    shift
-  else
-    exists() {
-      echo "$@" >&2
-      exit 1
-    }
-  fi
-
-  legend() {
-    if [ $quiet -eq 0 ]; then
-      echo "$@" >&2
-    fi
-  }
-
-  subcmd() {
-    local cmdname="''${1//-/_}"
-    shift
-    "subcmd_$cmdname" "$@"
-  }
-
-  subcmd_quiet() {
-    local prev_quiet=$quiet
-    quiet=1
-    subcmd "$@"
-    local ret=$?
-    quiet=$prev_quiet
-    return $ret
-  }
-
-  ${lib.concatStrings (lib.mapAttrsToList mkSubCommand subcommands)}
-
-  case "$1" in
-    ${lib.concatStrings (lib.mapAttrsToList mkCase subcommands)}
-    *) echo "Usage: ${commandName} <subcommand> [<args>]" >&2
-       echo >&2
-       echo "A tool to manage taskserver users on NixOS" >&2
-       echo >&2
-       echo "The following subcommands are available:" >&2
-       ${lib.concatMapStringsSep "\n     " (c: "usage_${mkShellName c}")
-                                           (lib.attrNames subcommands)}
-       exit 1
-  esac
-''
diff --git a/nixos/modules/services/misc/taskserver/helper-tool.py b/nixos/modules/services/misc/taskserver/helper-tool.py
new file mode 100644
index 000000000000..3277a50cd510
--- /dev/null
+++ b/nixos/modules/services/misc/taskserver/helper-tool.py
@@ -0,0 +1,276 @@
+import grp
+import pwd
+import os
+import re
+import string
+import subprocess
+import sys
+
+from shutil import rmtree
+from tempfile import NamedTemporaryFile
+
+import click
+
+CERTTOOL_COMMAND = "@certtool@"
+TASKD_COMMAND = "@taskd@"
+TASKD_DATA_DIR = "@dataDir@"
+TASKD_USER = "@user@"
+TASKD_GROUP = "@group@"
+FQDN = "@fqdn@"
+
+RE_CONFIGUSER = re.compile(r'^\s*user\s*=(.*)$')
+
+
+def run_as_taskd_user():
+    uid = pwd.getpwnam(TASKD_USER).pw_uid
+    gid = grp.getgrnam(TASKD_GROUP).gr_gid
+    os.setgid(gid)
+    os.setuid(uid)
+
+
+def taskd_cmd(cmd, *args, **kwargs):
+    return subprocess.call(
+        [TASKD_COMMAND, cmd, "--data", TASKD_DATA_DIR] + list(args),
+        preexec_fn=run_as_taskd_user,
+        **kwargs
+    )
+
+
+def label(msg):
+    if sys.stdout.isatty() or sys.stderr.isatty():
+        sys.stderr.write(msg + "\n")
+
+
+def mkpath(*args):
+    return os.path.join(TASKD_DATA_DIR, "orgs", *args)
+
+
+def fetch_username(org, key):
+    for line in open(mkpath(org, "users", key, "config"), "r"):
+        match = RE_CONFIGUSER.match(line)
+        if match is None:
+            continue
+        return match.group(1).strip()
+    return None
+
+
+def generate_key(org, user):
+    basedir = os.path.join(TASKD_DATA_DIR, "keys", org, user)
+    if os.path.exists(basedir):
+        raise OSError("Keyfile directory for {} already exists.".format(user))
+
+    privkey = os.path.join(basedir, "private.key")
+    pubcert = os.path.join(basedir, "public.cert")
+    cakey = os.path.join(TASKD_DATA_DIR, "keys", "ca.key")
+    cacert = os.path.join(TASKD_DATA_DIR, "keys", "ca.cert")
+
+    try:
+        os.makedirs(basedir, mode=0700)
+
+        cmd = [CERTTOOL_COMMAND, "-p", "--bits", "2048", "--outfile", privkey]
+        subprocess.call(cmd, preexec_fn=lambda: os.umask(0077))
+
+        template = NamedTemporaryFile(mode="w", prefix="certtool-template")
+        template.writelines(map(lambda l: l + "\n", [
+            "organization = {0}".format(org),
+            "cn = {}".format(FQDN),
+            "tls_www_client",
+            "encryption_key",
+            "signing_key"
+        ]))
+        template.flush()
+
+        cmd = [CERTTOOL_COMMAND, "-c",
+               "--load-privkey", privkey,
+               "--load-ca-privkey", cakey,
+               "--load-ca-certificate", cacert,
+               "--template", template.name,
+               "--outfile", pubcert]
+        subprocess.call(cmd, preexec_fn=lambda: os.umask(0077))
+    except:
+        rmtree(basedir)
+        raise
+
+
+def is_key_line(line, match):
+    return line.startswith("---") and line.lstrip("- ").startswith(match)
+
+
+def getkey(*args):
+    path = os.path.join(TASKD_DATA_DIR, "keys", *args)
+    buf = []
+    for line in open(path, "r"):
+        if len(buf) == 0:
+            if is_key_line(line, "BEGIN"):
+                buf.append(line)
+            continue
+
+        buf.append(line)
+
+        if is_key_line(line, "END"):
+            return ''.join(buf)
+    raise IOError("Unable to get key from {}.".format(path))
+
+
+def mktaskkey(cfg, path, keydata):
+    heredoc = 'cat > "{}" <<EOF\n{}EOF'.format(path, keydata)
+    cmd = 'task config taskd.{} -- "{}"'.format(cfg, path)
+    return heredoc + "\n" + cmd
+
+
+@click.group()
+@click.option('--service-helper', is_flag=True)
+@click.pass_context
+def cli(ctx, service_helper):
+    """
+    Manage Taskserver users and certificates
+    """
+    ctx.obj = {'is_service_helper': service_helper}
+
+
+@cli.command("list-users")
+@click.argument("organisation")
+def list_users(organisation):
+    """
+    List all users belonging to the specified organisation.
+    """
+    label("The following users exist for {}:".format(organisation))
+    for key in os.listdir(mkpath(organisation, "users")):
+        name = fetch_username(organisation, key)
+        if name is not None:
+            sys.stdout.write(name + "\n")
+
+
+@cli.command("list-orgs")
+def list_orgs():
+    """
+    List available organisations
+    """
+    label("The following organisations exist:")
+    for org in os.listdir(mkpath()):
+        sys.stdout.write(org + "\n")
+
+
+@cli.command("get-uuid")
+@click.argument("organisation")
+@click.argument("user")
+def get_uuid(organisation, user):
+    """
+    Get the UUID of the specified user belonging to the specified organisation.
+    """
+    for key in os.listdir(mkpath(organisation, "users")):
+        name = fetch_username(organisation, key)
+        if name is not None and name == user:
+            label("User {} has the following UUID:".format(name))
+            sys.stdout.write(key + "\n")
+            return
+    sys.exit("No UUID found for user {}.".format(user))
+
+
+@cli.command("export-user")
+@click.argument("organisation")
+@click.argument("user")
+def export_user(organisation, user):
+    """
+    Export user of the specified organisation as a series of shell commands
+    that can be used on the client side to easily import the certificates.
+
+    Note that the private key will be exported as well, so use this with care!
+    """
+    name = key = None
+    for current_key in os.listdir(mkpath(organisation, "users")):
+        name = fetch_username(organisation, current_key)
+        if name is not None and name == user:
+            key = current_key
+            break
+
+    if name is None:
+        msg = "User {} doesn't exist in organisation {}."
+        sys.exit(msg.format(user, organisation))
+
+    pubcert = getkey(organisation, user, "public.cert")
+    privkey = getkey(organisation, user, "private.key")
+    cacert = getkey("ca.cert")
+
+    keydir = "${TASKDATA:-$HOME/.task}/keys"
+
+    credentials = '/'.join([organisation, user, key])
+    allow_unquoted = string.ascii_letters + string.digits + "/-_."
+    if not all((c in allow_unquoted) for c in credentials):
+        credentials = "'" + credentials.replace("'", r"'\''") + "'"
+
+    script = [
+        "umask 0077",
+        'mkdir -p "{}"'.format(keydir),
+        mktaskkey("certificate", os.path.join(keydir, "public.cert"), pubcert),
+        mktaskkey("key", os.path.join(keydir, "private.key"), privkey),
+        mktaskkey("ca", os.path.join(keydir, "ca.cert"), cacert),
+        "task config taskd.credentials -- {}".format(credentials)
+    ]
+
+    sys.stdout.write('\n'.join(script))
+
+
+@cli.command("add-org")
+@click.argument("name")
+@click.pass_obj
+def add_org(obj, name):
+    """
+    Create an organisation with the specified name.
+    """
+    if os.path.exists(mkpath(name)):
+        if obj['is_service_helper']:
+            return
+        msg = "Organisation with name {} already exists."
+        sys.exit(msg.format(name))
+
+    taskd_cmd("add", "org", name)
+
+
+@cli.command("add-user")
+@click.argument("organisation")
+@click.argument("user")
+@click.pass_obj
+def add_user(obj, organisation, user):
+    """
+    Create a user for the given organisation along with a client certificate
+    and print the key of the new user.
+
+    The client certificate along with it's public key can be shown via the
+    'export-user' subcommand.
+    """
+    if not os.path.exists(mkpath(organisation)):
+        sys.exit("Organisation {} does not exist.".format(organisation))
+
+    if os.path.exists(mkpath(organisation, "users")):
+        for key in os.listdir(mkpath(organisation, "users")):
+            name = fetch_username(organisation, key)
+            if name is not None and name == user:
+                if obj['is_service_helper']:
+                    return
+                msg = "User {} already exists in organisation {}."
+                sys.exit(msg.format(user, organisation))
+
+    taskd_cmd("add", "user", organisation, user)
+    generate_key(organisation, user)
+
+
+@cli.command("add-group")
+@click.argument("organisation")
+@click.argument("group")
+@click.pass_obj
+def add_group(obj, organisation, group):
+    """
+    Create a group for the given organisation.
+    """
+    if os.path.exists(mkpath(organisation, "groups", group)):
+        if obj['is_service_helper']:
+            return
+        msg = "Group {} already exists in organisation {}."
+        sys.exit(msg.format(group, organisation))
+
+    taskd_cmd("add", "group", organisation, group)
+
+
+if __name__ == '__main__':
+    cli()