#!/usr/bin/python3
"""
Wrapper around cargo to have it build using Debian settings.

Usage:
    export PATH=/path/to/dir/of/this/script:$PATH
    export CARGO_HOME=debian/cargo_home
    cargo prepare-debian /path/to/local/registry [--link-from-system]
    cargo build
    cargo test
    cargo install
    cargo clean
    [rm -rf /path/to/local/registry]

The "prepare-debian" subcommand writes a config file to $CARGO_HOME that makes
the subsequent invocations use our Debian flags. The "--link-from-system" flag
is optional; if you use it we will create /path/to/local/registry and symlink
the contents of /usr/share/cargo/registry into it. You are then responsible for
cleaning it up afterwards (a simple `rm -rf` should do).

See cargo:d/rules and dh-cargo:cargo.pm for more examples.

Make sure you add "Build-Depends: python3:native" if you use this directly.
If using this only indirectly via dh-cargo, then you only need "Build-Depends:
dh-cargo"; this is a general principle when declaring dependencies.

If CARGO_HOME doesn't end with debian/cargo_home, then this script does nothing
and passes through directly to cargo.

Otherwise, this script builds in a debian-ish way.
As an Ubuntu change, it also attempts to wrap the "real" cargo invocation
in cargo-auditable, which embeds dependency metadata in any built binaries.
See: https://github.com/rust-secure-code/cargo-auditable
You *must* set the following environment variables:

- DEB_CARGO_CRATE
  ${crate}_${version} of whatever you're building.

- CFLAGS CXXFLAGS CPPFLAGS LDFLAGS [*]
- DEB_HOST_GNU_TYPE DEB_HOST_RUST_TYPE [*]

- (required only for `cargo install`) DESTDIR
  DESTDIR to install build artifacts under. If running via dh-cargo, this will
  be set automatically by debhelper, see `dh_auto_install` for details.

- (optional) DEB_BUILD_OPTIONS DEB_BUILD_PROFILES

- (optional) DEB_CARGO_INSTALL_PREFIX
  Prefix to install build artifacts under. Default: /usr. Sometimes you might
  want to change this to /usr/lib/cargo if the binary clashes with something
  else, and then symlink it into /usr/bin under an alternative name.

- (optional) DEB_CARGO_CRATE_IN_REGISTRY
  Whether the crate is in the local-registry (1) or cwd (0, empty, default).

- (optional) UBUNTU_ENABLE_CARGO_AUDITABLE
  If this is set to "1" or "true", wrap cargo invocations in cargo-auditable.
  This embeds dependency metadata into any built binaries.
  See d/README.debian

For the envvars marked [*], it is easiest to set these in your d/rules via:

    include /usr/share/dpkg/architecture.mk
    include /usr/share/dpkg/buildflags.mk
    include /usr/share/rustc/architecture.mk
    export CFLAGS CXXFLAGS CPPFLAGS LDFLAGS
    export DEB_HOST_RUST_TYPE DEB_HOST_GNU_TYPE
"""

import os
import os.path
import shutil
import subprocess
import sys
import json
import tomllib
import re
import textwrap

FLAGS = "CFLAGS CXXFLAGS CPPFLAGS LDFLAGS"
ARCHES = "DEB_HOST_GNU_TYPE DEB_HOST_RUST_TYPE"
SYSTEM_REGISTRY = "/usr/share/cargo/registry"

def log(*args):
    print("debian cargo wrapper:", *args, file=sys.stderr, flush=True)

def logrun(*args, **kwargs):
    log("running subprocess", args, kwargs)
    return subprocess.run(*args, **kwargs)

def sourcepath(p=None):
    return os.path.join(os.getcwd(), p) if p else os.getcwd()

def check_enable_cargo_auditable(cargo_home):
    enable_flag = os.getenv("UBUNTU_ENABLE_CARGO_AUDITABLE")
    try_enable = None
    if enable_flag == "1" or enable_flag == "true":
        try_enable = True
    elif enable_flag == "0" or enable_flag == "false":
        try_enable = False
    elif enable_flag is None:
        # Currently, we opt IN, so none == false
        log("UBUNTU_ENABLE_CARGO_AUDITABLE was unset, do not try to use it")
        try_enable = False
    else:
        raise ValueError(f"UBUNTU_ENABLE_CARGO_AUDITABLE should be 1, true, 0, or false (got {enable_flag!r})")

    if not enable_flag:
        return False

    cargo_audit_loc = "/usr/bin/cargo-auditable"
    if os.path.isfile(cargo_audit_loc):
        log(f"Found {cargo_audit_loc}, wrapping with cargo-auditable")
        return True
    else:
        raise ValueError(f"Requested UBUNTU_ENABLE_CARGO_AUDITABLE, but could not find {cargo_audit_loc}!")

# Return a list of three-tuples: `(name, version_requirement, [feature-flags])`
def get_dev_dependencies():
    # --no-deps means don't *examine* this package's dependencies.
    # it still outputs the dependencies of *this* package.
    cargo_metadata = logrun(
            ["/usr/lib/rust-1.93/bin/cargo", "metadata", "--no-deps", "--format-version=1", "--offline"],
            cwd=sourcepath(), stdout=subprocess.PIPE
    )
    metadata_json = cargo_metadata.stdout.decode('utf-8')
    metadata = json.loads(metadata_json)
    out = []
    for pkg in metadata["packages"]:
        for dep in pkg["dependencies"]:
            if dep["kind"] == "dev":
                out.append((dep["name"], dep["req"], dep["features"]))
    return out

def version_req_to_decl(version):
    # The version *requirement* of the top-level package does not need to be a full semver string,
    # but the version *declaration* of the fake package does.
    # [dependencies] libfoobar = "^1.0" --> [libfoobar] version = "1.0.0"
    # See here for the requirement syntax: https://doc.rust-lang.org/cargo/reference/specifying-dependencies.html#version-requirement-syntax
    # In general this is a really hard problem. I have tested this code on all of the ~40,000 version requirement strings
    # on crates.io, and it fails to generate a working string on 77 of them.
    # That's good enough for me. If it becomes a problem, use a d/patch to make the version req
    # less horrible.

    num_re = r"(\d+|~|\*)"
    # Match an allowed leading character, some spaces, then "num(.num(.num))"
    # Use ?: to avoid capturing the groups used to indicate "up to 3 dot separated components"
    whole_re = fr"(<|<=|=|>=|>|\^|~)?\s*{num_re}(?:\.{num_re}(?:\.{num_re})?)?"

    # only consider the first version string.
    # please, nobody make a req like ">=1.0, >=2.0"
    req = version.split(",")[0].strip()
    matched = re.match(whole_re, req)
    if matched is None:
        raise ValueError(f"invalid version req {version!r}")

    # remember, matched[0] is the entire match
    leader = matched[1]
    pre_info = matched.string[matched.end(0):]

    nums_out = []
    for i in range(3):
        # skip 0-index (entire match) and leader
        num = matched[i + 2]
        if num == "*" or num is None:
            nums_out.append(0)
        else:
            nums_out.append(int(num))
    # Usually, the version string indicated after the leader works as a dependency
    # (once missing components are filled in).
    # The edge cases are `>` and `<`
    # We don't need to do advanced checking to see if the version range requested is impossible
    # because Cargo would complain anyways.
    # What these mean precisely is largely undocumented, but it looks like
    # they want the rightmost non-zero component incremented or decremented
    if leader == ">" or leader == "<":
        # Find the index of the rightmost non-zero component
        for i in reversed(range(3)):
            if nums_out[i] != 0:
                twiddle = 1 if leader == ">" else -1
                nums_out[i] += twiddle
                break

    return ".".join(str(i) for i in nums_out) + pre_info

def create_dev_dependencies_stubs(cargo_config):
    # cargo-auditable uses cargo-metadata to get the list of crates used in the build,
    # but it doesn't filter out dev-dependencies currently.
    # When it tries to get the metadata for dev-deps, it fails because the package isn't there
    # (because it's only supposed to be there in dev configurations).
    # This is a Cargo-internal issue.
    # As a workaround, create fake stub packages.
    # See: https://github.com/rust-secure-code/cargo-auditable/issues/128
    # Thank you Shnatsel for all your help

    dev_deps = get_dev_dependencies()
    log(f"Found dev-dependencies: {dev_deps}")

    # Exfiltrate dh-cargo-deps path, as written to cargo config toml in prepare_debian
    # This way we know where to write the new package.
    with open(cargo_config, "rb") as f:
        toml = tomllib.load(f)
        registry_path = toml["source"]["dh-cargo-registry"]["directory"]

    for (pkg, ver, features) in dev_deps:
        full_version_string = version_req_to_decl(ver)
        fake_pkg_path = f"{registry_path}/{pkg}-{full_version_string}"
        log(f"have dev-dependency {pkg} = {ver}; checking '{fake_pkg_path}' to spoof dev dependency?")
        if os.path.exists(fake_pkg_path):
            log("^^ already exists, skipping")
        else:
            os.makedirs(fake_pkg_path, exist_ok=False)
            # a minimal package requires:
            # - Cargo.toml
            # - a src/lib.rs or src/main.rs
            # - a .cargo-checksum.json, although it can be empty --
            #   apparently it will be required to be populated at some point,
            #   but I am hoping that native SBOM support ships before then ;')

            # Cargo doesn't care about `+etc` versions, so append this to make it very clear
            # that this is not a real package
            hint_version_string = full_version_string + "+ubuntu-stub"
            # The native Python toml library doesn't support writing,
            # so I have to format the toml myself.
            package_section = textwrap.dedent(f"""\
                [package]
                name = {repr(pkg)}
                version = {repr(hint_version_string)}
                description = "[UBUNTU] Stub version of {pkg} to spoof dev dependency" """)
            # [features]
            # foo = []
            # bar = []
            features_entries = "\n".join(f"{feature_name} = []" for feature_name in features)
            with open(fake_pkg_path + "/Cargo.toml", "w") as f:
                f.write(package_section + "\n\n[features]\n" + features_entries)

            os.makedirs(fake_pkg_path + "/src")
            open(fake_pkg_path + "/src/lib.rs", "x").close()
            with open(fake_pkg_path + "/.cargo-checksum.json", "x") as f:
                f.write('{ "files": {} }')
            log(f"^^ created stub package {pkg} @ {full_version_string}")

def prepare_debian(cargo_home, registry, cratespec, host_gnu_type, ldflags, link_from_system, extra_rustflags):
    registry_path = sourcepath(registry)
    if link_from_system:
        log("linking %s/* into %s/" % (SYSTEM_REGISTRY, registry_path))
        os.makedirs(registry_path, exist_ok=True)
        crates = os.listdir(SYSTEM_REGISTRY) if os.path.isdir(SYSTEM_REGISTRY) else []
        for c in crates:
            target = os.path.join(registry_path, c)
            if not os.path.islink(target):
                os.symlink(os.path.join(SYSTEM_REGISTRY, c), target)
    elif not os.path.exists(registry_path):
        raise ValueError("non-existent registry: %s" % registry)

    rustflags = "-C debuginfo=2 -C strip=none --cap-lints warn".split()
    rustflags.extend(["-C", "linker=%s-gcc" % host_gnu_type])
    for f in ldflags:
        rustflags.extend(["-C", "link-arg=%s" % f])
    if link_from_system:
        rustflags.extend([
            # Note that this order is important! Rust evaluates these options in
            # priority of reverse order, so if the second option were in front,
            # it would never be used, because any paths in registry_path are
            # also in sourcepath().
            "--remap-path-prefix", "%s=%s/%s" %
                (sourcepath(), SYSTEM_REGISTRY, cratespec.replace("_", "-")),
            "--remap-path-prefix", "%s=%s" % (registry_path, SYSTEM_REGISTRY),
        ])
    rustflags.extend(extra_rustflags.split())

    # TODO: we cannot enable this until dh_shlibdeps works correctly; atm we get:
    # dpkg-shlibdeps: warning: can't extract name and version from library name 'libstd-XXXXXXXX.so'
    # and the resulting cargo.deb does not depend on the correct version of libstd-rust-1.XX
    # We probably need to add override_dh_makeshlibs to d/rules of rustc
    #rustflags.extend(["-C", "prefer-dynamic"])

    os.makedirs(cargo_home, exist_ok=True)
    with open("%s/config.toml" % cargo_home, "w") as fp:
        fp.write("""[source.crates-io]
replace-with = "dh-cargo-registry"

[source.dh-cargo-registry]
directory = "{0}"

[build]
rustflags = {1}

[profile.release]
debug = true
""".format(registry_path, repr(rustflags)))

    return 0

def install(destdir, cratespec, host_rust_type, crate_in_registry, install_prefix, *args):
    crate, version = cratespec.rsplit("_", 1)
    log("installing into destdir '%s' prefix '%s'" % (destdir, install_prefix))
    install_target = destdir + install_prefix
    logrun(["env", "RUST_BACKTRACE=1",
        # set CARGO_TARGET_DIR so build products are saved in target/
        # normally `cargo install` deletes them when it exits
        "CARGO_TARGET_DIR=" + sourcepath("target"),
        "/usr/bin/cargo"] + list(args) +
        ([crate, "--vers", version] if crate_in_registry else ["--path", sourcepath()]) +
        ["--root", install_target], check=True)
    logrun(["rm", "-f", "%s/.crates.toml" % install_target])
    logrun(["rm", "-f", "%s/.crates2.json" % install_target])

    # if there was a custom build output, symlink it to debian/cargo_out_dir
    # hopefully cargo will provide a better solution in future https://github.com/rust-lang/cargo/issues/5457
    r = logrun('''ls -td "target/%s/release/build/%s"-*/out 2>/dev/null | head -n1'''
        % (host_rust_type, crate), shell=True, stdout=subprocess.PIPE).stdout
    r = r.decode("utf-8").rstrip()
    if r:
        logrun(["ln", "-sfT", "../%s" % r, "debian/cargo_out_dir"], check=True)
    return 0

def main(*args):
    cargo_home = os.getenv("CARGO_HOME", "")
    if not cargo_home.endswith("/debian/cargo_home"):
        os.execv("/usr/lib/rust-1.93/bin/cargo", ["cargo"] + list(args))

    if any(f not in os.environ for f in FLAGS.split()):
        raise ValueError("not all of %s set; did you call dpkg-buildflags?" % FLAGS)

    if any(f not in os.environ for f in ARCHES.split()):
        raise ValueError("not all of %s set; did you include architecture.mk?" % ARCHES)

    build_options = os.getenv("DEB_BUILD_OPTIONS", "").split()
    build_profiles = os.getenv("DEB_BUILD_PROFILES", "").split()

    enable_cargo_auditable = check_enable_cargo_auditable(cargo_home)

    parallel = []
    lto = 0
    for o in build_options:
        if o.startswith("parallel="):
            parallel = ["-j" + o[9:]]
        elif o.startswith("optimize="):
            opt_arg = o[9:]
            for arg in opt_arg.split(","):
                if opt_arg == "-lto":
                    lto = -1
                elif opt_arg == "+lto":
                    lto = 1
                else:
                    log(f"WARNING: unhandled optimization flag: {opt_arg}")

    nodoc = "nodoc" in build_options or "nodoc" in build_profiles
    nocheck = "nocheck" in build_options or "nocheck" in build_profiles

    # note this is actually the "build target" type, see rustc's README.Debian
    # for full details of the messed-up terminology here
    host_rust_type = os.getenv("DEB_HOST_RUST_TYPE", "")
    host_gnu_type = os.getenv("DEB_HOST_GNU_TYPE", "")

    log("options, profiles, parallel, lto:", build_options, build_profiles, parallel, lto)
    log("rust_type, gnu_type:", ", ".join([host_rust_type, host_gnu_type]))

    if "RUSTFLAGS" in os.environ:
        # see https://github.com/rust-lang/cargo/issues/6338 for explanation on why we must do this
        log("unsetting RUSTFLAGS and assuming it will be (or already was) added to $CARGO_HOME/config.toml")
        extra_rustflags = os.environ["RUSTFLAGS"]
        del os.environ["RUSTFLAGS"]
    else:
        extra_rustflags = ""

    if args[0] == "prepare-debian":
        registry = args[1]
        link_from_system = False
        if len(args) > 2 and args[2] == "--link-from-system":
            link_from_system = True
        return prepare_debian(cargo_home, registry,
            os.environ["DEB_CARGO_CRATE"], host_gnu_type,
            os.getenv("LDFLAGS", "").split(), link_from_system, extra_rustflags)

    newargs = []
    if enable_cargo_auditable:
        newargs.append("auditable")
    subcmd = None
    for a in args:
        if (subcmd is None) and (a in ("build", "rustc", "doc", "test", "bench", "install")):
            subcmd = a
            newargs.extend(["-Zavoid-dev-deps", a, "--verbose", "--verbose"] +
                parallel + ["--target", host_rust_type])
        elif (subcmd is None) and (a == "clean"):
            subcmd = a
            newargs.extend([a, "--verbose", "--verbose"])
        else:
            newargs.append(a)

    if subcmd is not None and "--verbose" in newargs and "--quiet" in newargs:
        newargs.remove("--quiet")

    if nodoc and subcmd == "doc":
        return 0
    if nocheck and subcmd in ("test", "bench"):
        return 0

    if lto == 1:
        newargs.extend(["--config", "profile.release.lto = \"thin\""])
    elif lto == -1:
        newargs.extend(["--config", "profile.release.lto = false"])

    if subcmd == "clean":
        logrun(["env", "RUST_BACKTRACE=1", "/usr/bin/cargo"] + list(newargs), check=True)
        if os.path.exists(cargo_home):
            shutil.rmtree(cargo_home)
        return 0

    cargo_config = "%s/config.toml" % cargo_home
    if not os.path.exists(cargo_config):
        raise ValueError("does not exist: %s, did you run `cargo prepare-debian <registry>`?" % cargo_config)

    if enable_cargo_auditable and (subcmd == "build" or subcmd == "install"):
        create_dev_dependencies_stubs(cargo_config)

    if subcmd == "install":
        return install(os.getenv("DESTDIR", ""),
            os.environ["DEB_CARGO_CRATE"],
            host_rust_type,
            os.getenv("DEB_CARGO_CRATE_IN_REGISTRY", "") == "1",
            os.getenv("DEB_CARGO_INSTALL_PREFIX", "/usr"),
            *newargs)
    else:
        return logrun(["env", "RUST_BACKTRACE=1", "/usr/bin/cargo"] + list(newargs)).returncode

if __name__ == "__main__":
    sys.exit(main(*sys.argv[1:]))
