manage-dotfiles

#!/usr/bin/python3

from __future__ import annotations

import argparse
import hashlib
import os
import socket
import stat
import subprocess as subp
import sys
import tempfile
from pathlib import Path


ENCRYPTED_FILES = {
    ".atlas/auth": "atlas-auth.gpg",
    ".gitspindle": "gitspindle.gpg",
    ".gnupg/gpg.conf": "gpg.conf.gpg",
    ".gnupg/gpg-agent.conf": "gpg-agent.conf.gpg",
    ".netrc": "netrc.gpg",
    ".npmrc": "npmrc.gpg",
    ".ssh/authorized_keys": "ssh-authorized-keys.gpg",
    ".ssh/config": "ssh-config.gpg",
    ".transifexrc": "transifexrc.gpg",
    ".config/aider/env": "aider/env.gpg",
    ".config/git/doist.inc": "git/doist.inc.gpg",
    ".config/exist-git-hook.conf": "exist-git-hook.conf.gpg",
    ".config/isyncrc": "isyncrc.gpg",
    ".config/pws.yaml": "pws.yaml.gpg",
    ".config/emacs/init-99-private.el": "emacs/init-99-private.el.gpg",
    ".config/fish/conf.d/doist.fish": "fish/conf.d/doist.fish.gpg",
    ".config/fish/conf.d/private.fish": "fish/conf.d/private.fish.gpg",
    ".config/mpdasrc": "mpdasrc.gpg",
    ".config/msmtp/config": "msmtp/config.gpg",
    ".config/notmuch/autotag.py": "notmuch/autotag.py.gpg",
    ".config/notmuch/default/config": "notmuch/default/config.gpg",
    ".config/spotify-player/app.toml.in": "spotify-player/app.toml.in.gpg",
    ".config/swww/choose-folder": "swww/choose-folder.gpg",
}

ENCRYPTED_TARBALLS = {
    "Age": ("age.tar.gpg", ".config/age/*"),
    "Calibre": ("calibre.tar.gpg", ".config/calibre/*"),
    "llmcli": ("llmcli.tar.gpg", ".config/llmcli/*"),
    "Pgweb": ("pgweb.tar.gpg", ".pgweb/**/**"),
    "Systemd mounts": ("systemd-mounts.tar.gpg", ".config/systemd/user/*.mount"),
    "Ulauncher": (
        "ulauncher.tar.gpg",
        ".config/ulauncher/*",
        ".config/ulauncher/*/*",
    ),
    "Weechat": ("weechat.tar.gpg", ".config/weechat/*.conf"),
}

FILE_SYMLINKS = {
    "aider/config.yml": ".aider.conf.yml",
    "kanshi/config.@@HOSTNAME@@": ".config/kanshi/config",
    "notmuch/default/config": ".notmuch-config",
}

DIR_SYMLINKS = {
    "applications": ".local/share/applications",
    "git/config": ".gitconfig",
    "xkb": ".xkb",
}

EXECUTABLE_FILES = {
    ".config/notmuch/autotag.py",
    ".config/swww/choose-folder",
}

TEMPLATED_FILES = {
    ".config/spotify-player/app.toml",
}

HOME = Path.home()
CONFIG = HOME / ".config"

_filters = {}


def filter(filename: str):
    def wrapper(func):
        _filters[filename] = func
        return func

    return wrapper


class TempPath(os.PathLike):
    def __init__(self):
        self.tmpf = tempfile.NamedTemporaryFile()

    def __fspath__(self):
        return self.tmpf.name

    @property
    def path(self):
        return Path(self.tmpf.name)


@filter(".netrc")
def filter_netrc(fp: Path) -> os.PathLike:
    data = fp.read_text()
    new_data = []
    machine, login = "", ""
    for line in data.splitlines():
        if not line.strip():
            continue
        words = line.strip().split()
        if words[0] == "machine":
            machine, login = words[1], ""
        elif words[0] == "login":
            login = words[1]
        elif (
            words[0] == "password"
            and machine.endswith(".amazonaws.com")
            and login == "aws"
        ):
            line = "\tpassword ***"
        new_data.append(line)
    outp = TempPath()
    outp.path.write_text("\n".join(new_data))
    return outp


def render_template(src_fp: Path, dst_fp: Path) -> None:
    hostname = socket.gethostname()
    replacements = {
        "@@HOSTNAME@@": hostname,
        "@@HOSTNAME_TITLE@@": hostname.title(),
    }

    text = src_fp.read_text()
    for needle, value in replacements.items():
        text = text.replace(needle, value)
    dst_fp.write_text(text)


def md5(buf: bytes) -> str:
    hh = hashlib.md5()
    hh.update(buf)
    return hh.hexdigest()


def get_last_update_filename(name: str) -> str:
    dst_name = CONFIG / name
    dst_hash = md5(str(dst_name).encode() + b"\n")
    return f".last_update_{dst_hash}"


def _age_recipients():
    age_dir = CONFIG / "age"
    rcpts = []
    for pub_file in age_dir.glob("*.pub"):
        rcpts += ["-R", pub_file]
    return rcpts


def _age_identities():
    age_dir = CONFIG / "age"
    ids = []
    for id_file in age_dir.glob("*identity*.txt"):
        ids += ["-i", id_file]
    return ids


def encrypt_data(src: Path | bytes, dst_p: Path):
    if dst_p.suffix == ".gpg":
        cmd = ["gpg2", "--quiet", "--batch", "--output", dst_p, "--encrypt"]
    elif dst_p.suffix == ".age":
        cmd = ["age", "--encrypt"] + _age_recipients() + ["-o", dst_p]
    else:
        raise ValueError("Unknown suffix for encrypted file: " + dst_p.suffix)
    input_ = src
    if isinstance(src, Path):
        cmd.append(src)
        input_ = None
    return subp.run(cmd, check=True, input=input_)


def decrypt_data(src_p: Path, dst: Path | None = None):
    if src_p.suffix == ".gpg":
        cmd = ["gpg2", "--quiet", "--batch", "--decrypt", src_p]
    elif src_p.suffix == ".age":
        cmd = ["age", "--decrypt"] + _age_identities() + [src_p]
    else:
        raise ValueError("Unknown suffix for encrypted file: " + src_p.suffix)
    capture_output = True
    if dst is not None:
        cmd = cmd[:-1] + ["--output", dst, src_p]
        capture_output = False
    res = subp.run(cmd, check=True, capture_output=capture_output)
    if capture_output:
        return res.stdout


def get_wanted_permissions(src: Path) -> int:
    with src.open("rb") as fd:
        data = fd.read(64)
    if data.startswith(b"#!"):
        return 0o700
    return 0o600


def update_encrypted_files(force: bool, verbose: bool):
    for src, dst in ENCRYPTED_FILES.items():
        src_p = HOME / src
        dst_p = CONFIG / dst
        uf_p = CONFIG / get_last_update_filename(dst)

        if verbose:
            print(f"Checking {dst_p} using {src_p}")

        if not src_p.exists():
            print(f"Missing source file: {src}")
            continue

        if (
            force
            or not dst_p.exists()
            or not uf_p.exists()
            or src_p.stat().st_mtime > uf_p.stat().st_mtime
        ):
            print(f"New version of {src} found, encrypting it...")

            if filter_func := _filters.get(src):
                print(f"  applying {filter_func.__name__}...")
                src_p = filter_func(src_p)

            # Check if content has changed
            if dst_p.exists() and decrypt_data(dst_p) == Path(src_p).read_bytes():
                print(f"  content is identical, skipping it...")
                uf_p.touch()
                continue

            dst_p.unlink(missing_ok=True)
            encrypt_data(Path(src_p), dst_p)
            uf_p.touch()


def update_encrypted_tarballs(force: bool, verbose: bool):
    for name, (dst, *patterns) in ENCRYPTED_TARBALLS.items():
        src_ps = []
        for pattern in patterns:
            src_ps += list(HOME.glob(pattern))
        src_ps = sorted(src_ps)
        dst_p = CONFIG / dst
        uf_p = CONFIG / get_last_update_filename(dst)

        if verbose:
            print(f"Checking {name} tarball ({dst_p}) using:")
            for p in src_ps:
                print(f"  {p.relative_to(HOME)}")

        # Inspired by https://github.com/maxmcd/reptar
        tar_cmd = [
            "tar",
            "--sort=name",
            "--mtime='1970-01-01 00:00:00Z'",
            "-c",
            "-C",
            HOME,
        ]
        tar_cmd += [p.relative_to(HOME) for p in src_ps]
        tar_data = subp.run(tar_cmd, check=True, capture_output=True).stdout
        tar_hh = md5(tar_data)
        uf_hh = "" if not uf_p.exists() else uf_p.read_text()

        if force or not dst_p.exists() or not uf_p.exists() or tar_hh != uf_hh:
            print(f"New version of {name} found, encrypting it...")
            dst_p.unlink(missing_ok=True)
            encrypt_data(tar_data, dst_p)
            uf_p.write_text(tar_hh)


def create_dir_symlinks():
    success = True

    for src, dst in DIR_SYMLINKS.items():
        src_p = CONFIG / src
        dst_p = HOME / dst

        if not dst_p.exists():
            print(f"Creating symlink: ~/{dst}")
            dst_p.symlink_to(src_p, target_is_directory=True)
        elif dst_p.is_symlink():
            actual_target = dst_p.resolve()
            if actual_target != src_p:
                print(
                    f"~/{dst} points to a wrong target :( ({actual_target}; should be {src})"
                )
                success = False
        elif dst_p.is_dir():
            if len(list(dst_p.iterdir())) == 0:
                print(f"Replace empty dir with symlink: ~/{dst}")
                dst_p.rmdir()
                dst_p.symlink_to(src_p, target_is_directory=True)
            else:
                print(f"~/{dst} exists and is not empty :(")
                success = False
        else:
            print(f"There is something wrong with ~/{dst} :/")
            success = False

    return success


def create_file_symlinks():
    success = True

    hostname = socket.gethostname()
    for src, dst in FILE_SYMLINKS.items():
        src = src.replace("@@HOSTNAME@@", hostname)
        src_p = CONFIG / src
        dst_p = HOME / dst

        if not dst_p.exists():
            print(f"Creating symlink: ~/{dst}")
            dst_p.symlink_to(src_p)
        elif dst_p.is_symlink():
            actual_target = dst_p.resolve()
            if actual_target != src_p:
                print(
                    f"~/{dst} points to a wrong target :( ({actual_target}; should be {src})"
                )
                success = False
        else:
            print(f"There is something wrong with ~/{dst} :/")
            success = False

    return success


def decrypt_files():
    for dst, src in ENCRYPTED_FILES.items():
        src_p = CONFIG / src
        dst_p = HOME / dst
        uf_p = CONFIG / get_last_update_filename(src)

        if (
            not dst_p.exists()
            or not uf_p.exists()
            or src_p.stat().st_mtime > uf_p.stat().st_mtime
        ):
            print(f"Decrypting ~/{dst}")
            dst_p.unlink(missing_ok=True)
            decrypt_data(src_p, dst_p)
            perms = get_wanted_permissions(dst_p)
            dst_p.lchmod(perms)
            uf_p.touch()

    return True


def decrypt_tarballs():
    for name, (src, *_) in ENCRYPTED_TARBALLS.items():
        src_p = CONFIG / src
        uf_p = CONFIG / get_last_update_filename(src)

        if not uf_p.exists() or src_p.stat().st_mtime > uf_p.stat().st_mtime:
            print(f"Decrypting {name} tarball")
            tar_data = decrypt_data(src_p)
            tar_hh = md5(tar_data)

            subp.run(["tar", "x", "-m", "-C", HOME], check=True, input=tar_data)
            uf_p.write_text(tar_hh)

    return True


def render_templates():
    success = True

    for path in TEMPLATED_FILES:
        dst_p = HOME / path
        src_p = dst_p.with_suffix(dst_p.suffix + ".in")
        if not src_p.exists():
            print(f"Template missing: {src_p}")
            success = False
            continue
        if not dst_p.exists() or src_p.stat().st_mtime > dst_p.stat().st_mtime:
            print(f"Rendering {dst_p}")
            render_template(src_p, dst_p)

    return success


def update_executable_files():
    perm = stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH
    for src in EXECUTABLE_FILES:
        src_p = HOME / src
        if not src_p.exists():
            continue
        mode = src_p.stat().st_mode
        wanted_mode = mode | stat.S_IXUSR
        if mode & stat.S_IRGRP:
            wanted_mode |= stat.S_IXGRP
        if mode & stat.S_IROTH:
            wanted_mode |= stat.S_IXOTH
        if mode != wanted_mode:
            src_p.chmod(wanted_mode)
    return True


def cmd_install(args):
    hook_path = Path(".git/hooks/pre-commit")
    if hook_path.exists() and not args.force:
        print("A pre-commit hook is already installed!")
        sys.exit(1)

    script_path = Path(__file__)

    script_content = f"""#!/usr/bin/bash

exec "{script_path}" encrypt "$@"
"""
    hook_path.write_text(script_content)


def cmd_encrypt(args):
    update_encrypted_files(args.force, args.verbose)
    update_encrypted_tarballs(args.force, args.verbose)


def cmd_list(args):
    dir_names = {p.name for p in CONFIG.glob(".last_update_*")}
    known_names = {
        get_last_update_filename(dst): src for src, dst in ENCRYPTED_FILES.items()
    }
    known_names.update(
        {
            get_last_update_filename(data[0]): f"{name} tarball"
            for name, data in ENCRYPTED_TARBALLS.items()
        }
    )
    known_names_set = set(known_names.keys())

    present_names = dir_names & known_names_set
    missing_names = known_names_set - dir_names
    unknown_names = dir_names - known_names_set

    if present_names:
        print("PRESENT UPDATE FILES:")
        for name in present_names:
            print(f"{known_names[name]}: {name}")
        print()
    if missing_names:
        print("MISSING UPDATE FILES:")
        for name in missing_names:
            print(f"{known_names[name]}: {name}")
        print()
    if unknown_names:
        print("UNKNOWN UPDATE FILES:")
        for name in unknown_names:
            print(name)


def cmd_setup(args):
    r1 = create_dir_symlinks()
    r2 = create_file_symlinks()
    r3 = decrypt_files()
    r4 = decrypt_tarballs()
    r5 = render_templates()
    r6 = update_executable_files()

    success = r1 and r2 and r3 and r4 and r5 and r6
    sys.exit(0 if success else 1)


if __name__ == "__main__":
    parser = argparse.ArgumentParser("Dotfiles pre-commit hook")
    sp = parser.add_subparsers()

    parser_install = sp.add_parser("install", help="install pre-commit hook")
    parser_install.add_argument(
        "-f", "--force", action="store_true", help="Allow overwriting existing files"
    )
    parser_install.set_defaults(func=cmd_install)

    parser_list = sp.add_parser("list", help="list update files")
    parser_list.set_defaults(func=cmd_list)

    parser_encrypt = sp.add_parser("encrypt", help="update encrypted files")
    parser_encrypt.add_argument(
        "-f", "--force", action="store_true", help="Force update"
    )
    parser_encrypt.add_argument(
        "-v", "--verbose", action="store_true", help="Verbose mode"
    )
    parser_encrypt.set_defaults(func=cmd_encrypt)

    parser_setup = sp.add_parser(
        "setup", help="setup symlinks and decrypt encrypted files"
    )
    parser_setup.set_defaults(func=cmd_setup)

    args = parser.parse_args()

    if "func" in args:
        args.func(args)
    else:
        parser.print_help()
        sys.exit(1)