#!/usr/bin/python3 -I
# wren-watch: written by Wren Linux parts/startup-watch.sh, installed as
# /usr/local/bin/wren-watch (root's). Watches the files a program running as
# you would change to start again at every login (your shell's startup files,
# sway's config, autostart, systemd user units, ~/.local/bin, ...): their
# fingerprints are kept by root, in /var/lib/wren-watch, where nothing running
# as you can change them, and a root timer compares them every 15 minutes and
# at boot. What changed shows in the bar ("startup changed") and here.
#   wren-watch                    what changed since you last accepted
#   wren-watch diff               the changes themselves (text files); asks sudo
#   wren-watch accept [FILE...]   take them as yours (all, or those files); asks sudo
#   wren-watch manifest           (for wrenctl) the fingerprints of yours, now
#   wren-watch check-all          (for root's timer) compare everyone watched
# It notices, it doesn't prevent: a program running as you can still change
# the files; it can't change the record, or keep the bar from showing it
# (sway-status is in ~/.local/bin, which is watched too).

# Python, standard library only, started with -I (no PYTHON* settings, no user site, nothing
# imported from the folder it is in). Root runs it on folders a program running as you
# controls, so every folder is opened from the one above it, never by its path, and never
# through a link: one swapped for a link half way leads nowhere.
import errno
import hashlib
import os
import pwd
import re
import shlex
import shutil
import stat
import subprocess
import sys

os.umask(0o022)
os.environ["PATH"] = "/usr/sbin:/usr/bin:/sbin:/bin"
# What sudo runs when diff or accept asks for it: root's copy, by its path,
# never one a program running as you could have put first in PATH
SELF = "/usr/local/bin/wren-watch"
# (a test's own pretend sudo: this runs as you then, so it crosses nothing;
# sudo clears it for root's run)
SUDO = os.environ.get("WREN_WATCH_SUDO") or "sudo"
# (a test points these at folders of its own)
TEST = bool(os.environ.get("WREN_WATCH_STATE"))
STATE = os.environ.get("WREN_WATCH_STATE") or "/var/lib/wren-watch"
RUN = os.environ.get("WREN_WATCH_RUN") or "/run/wren-watch"

FILES = [".profile", ".bash_profile", ".bash_login", ".bashrc", ".bash_logout", ".pam_environment",
         ".xprofile", ".xsessionrc", ".zshrc", ".zprofile", ".zshenv", ".gitconfig", ".config/git/config",
         ".config/mimeapps.list", ".ssh/rc", ".ssh/config", ".ssh/authorized_keys",
         ".gnupg/gpg-agent.conf", ".gnupg/gpg.conf", ".config/wren/autostart",
         ".tmux.conf", ".config/tmux/tmux.conf", ".vimrc", ".gdbinit", ".inputrc", ".wgetrc", ".curlrc",
         ".bash_aliases", ".bash_completion", ".config/foot/foot.ini"]
DIRS = [".config/sway", ".config/autostart", ".config/systemd/user", ".local/share/systemd/user", ".config/environment.d",
        ".local/bin", ".config/mako", ".local/share/applications", ".config/wren/on-lock.d",
        ".local/share/bash-completion", ".config/bash_completion.d"]
# Read at most this much of each file, and copy for diff only text files up to COPY_MAX
READ_MAX, COPY_MAX = 64 << 20, 256 << 10
# Folders deeper than this inside a watched one are not walked: one entry, "deep", stands for
# them (nothing starts from that far down, and a tree made deep on purpose can't stop the look)
DEPTH_MAX = 64
DIR_FLAGS = os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW | os.O_CLOEXEC


def die(msg):
    print("wren-watch: " + msg, file=sys.stderr)
    sys.exit(1)


# Control characters (and bytes that aren't UTF-8) never reach your terminal: a file's name and
# what is in it are a program's choice, and an escape sequence would be run by the terminal
_CNTRL = re.compile("[\x00-\x08\x0b-\x1f\x7f-\x9f\udc80-\udcff]")


def shown(text):
    return _CNTRL.sub("?", text)


def say(text=""):
    sys.stdout.write(shown(text) + "\n")


def name_of(rel):
    """How a path is written in the record and the list: as it is, or, with a newline, a
    control character or bytes that aren't UTF-8 in it, Python's escaped form of it (one line
    of plain text, which the bar, wrenctl and your terminal can take)."""
    if "\n" in rel or _CNTRL.search(rel):
        return repr(rel)
    return rel


def read_lines(path):
    try:
        with open(path, encoding="utf-8", errors="surrogateescape") as f:
            return [line.rstrip("\n") for line in f]
    except FileNotFoundError:
        return []


def write_file(path, text, mode):
    """PATH, in full or not at all: written next to it, then put in its place"""
    new = path + ".new"
    try:
        os.unlink(new)
    except FileNotFoundError:
        pass
    fd = os.open(new, os.O_WRONLY | os.O_CREAT | os.O_EXCL | os.O_NOFOLLOW | os.O_CLOEXEC, mode)
    with os.fdopen(fd, "w", encoding="utf-8", errors="surrogateescape") as f:
        os.fchmod(f.fileno(), mode)
        f.write(text)
    os.replace(new, path)
    return path


def remove(path):
    try:
        os.unlink(path)
    except FileNotFoundError:
        pass


# ---- the manifest: one line per watched thing under HOME, "KIND MODE SIZE HASH PATH"
# (PATH relative to HOME, last), without following any link, reading regular files only and
# at most READ_MAX of each; with COPIES, a copy of every text file up to COPY_MAX goes there,
# for diff

class Manifest:
    def __init__(self, home, copies=None):
        self.copies = copies
        self.out = []
        try:
            # (HOME itself is root's word, from the user list: the folders in it are not)
            self.home = os.open(home, os.O_RDONLY | os.O_DIRECTORY | os.O_CLOEXEC)
        except OSError:
            self.home = None

    def add(self, kind, mode, size, digest, rel):
        self.out.append("%s %s %s %s %s" % (kind, mode, size, digest, name_of(rel)))

    def parent(self, rel):
        """The folder REL is in, opened from HOME one step at a time; "linked" when a link is on
        the way (~/.ssh pointing at root's, say: root reads these files, so what is behind one
        is neither read nor copied); None when it isn't there"""
        fd = os.dup(self.home)
        try:
            for part in rel.split("/")[:-1]:
                try:
                    st = os.stat(part, dir_fd=fd, follow_symlinks=False)
                except OSError:
                    return None
                if stat.S_ISLNK(st.st_mode):
                    return "linked"
                if not stat.S_ISDIR(st.st_mode):
                    return None
                try:
                    nfd = os.open(part, DIR_FLAGS, dir_fd=fd)
                except OSError:
                    # it was a folder a moment ago: swapped since, so counted as a link
                    return "linked"
                os.close(fd)
                fd = nfd
            result, fd = fd, None
            return result
        finally:
            if fd is not None:
                os.close(fd)

    def entry(self, dirfd, name, rel):
        """NAME in the open folder DIRFD, as REL; its lstat, or None when it isn't there"""
        try:
            st = os.stat(name, dir_fd=dirfd, follow_symlinks=False)
        except OSError:
            return None
        mode = "%o" % stat.S_IMODE(st.st_mode)
        try:
            if stat.S_ISLNK(st.st_mode):
                target = os.readlink(name, dir_fd=dirfd)
                self.add("link", mode, "0", hashlib.sha256(os.fsencode(target)).hexdigest(), rel)
            elif stat.S_ISREG(st.st_mode):
                self.regular(dirfd, name, rel, st, mode)
            elif stat.S_ISDIR(st.st_mode):
                self.add("dir", mode, "0", "-", rel)
            else:
                self.add("other", mode, "0", "-", rel)
        except OSError:
            # changed under it (a link gone between two looks, say): a line unlike any other,
            # so it shows as changed, and the look goes on
            self.add("other", mode, "0", "changing", rel)
        return st

    def regular(self, dirfd, name, rel, st, mode):
        try:
            fd = os.open(name, os.O_RDONLY | os.O_NOFOLLOW | os.O_NONBLOCK | os.O_CLOEXEC, dir_fd=dirfd)
        except OSError:
            self.add("file", mode, str(st.st_size), "unreadable", rel)
            return
        with os.fdopen(fd, "rb") as f:
            if not stat.S_ISREG(os.fstat(f.fileno()).st_mode):
                # swapped for a pipe or such since the look above: never read
                self.add("other", mode, "0", "-", rel)
                return
            h, left, data = hashlib.sha256(), READ_MAX, b""
            while left > 0:
                b = f.read(min(1 << 20, left))
                if not b:
                    break
                h.update(b)
                left -= len(b)
                if self.copies and len(data) <= COPY_MAX:
                    data += b
        self.add("file", mode, str(st.st_size), h.hexdigest(), rel)
        # (only under a plain name: the copies are found again by the name the list gives)
        if self.copies and len(data) <= COPY_MAX and b"\0" not in data and name_of(rel) == rel:
            dest = os.path.join(self.copies, rel)
            os.makedirs(os.path.dirname(dest), mode=0o700, exist_ok=True)
            with open(os.open(dest, os.O_WRONLY | os.O_CREAT | os.O_TRUNC | os.O_NOFOLLOW, 0o600), "wb") as c:
                c.write(data)

    def walk(self, dirfd, rel, depth):
        """Everything in the open folder DIRFD (REL), and in the folders in it, in name order"""
        if depth > DEPTH_MAX:
            self.add("deep", "0", "0", "-", rel)
            return
        try:
            names = sorted(os.listdir(dirfd))
        except OSError:
            return
        folders = []
        for name in names:
            st = self.entry(dirfd, name, rel + "/" + name)
            if st is not None and stat.S_ISDIR(st.st_mode):
                folders.append(name)
        for name in folders:
            try:
                fd = os.open(name, DIR_FLAGS, dir_fd=dirfd)
            except OSError:
                continue
            try:
                self.walk(fd, rel + "/" + name, depth + 1)
            finally:
                os.close(fd)

    def lines(self):
        if self.home is None:
            return []
        for rel in FILES + DIRS:
            fd = self.parent(rel)
            if fd == "linked":
                self.add("linked", "0", "0", "-", rel)
                continue
            if fd is None:
                continue
            name = rel.rsplit("/", 1)[-1]
            try:
                st = self.entry(fd, name, rel)
                # a watched folder: what is in it too (one that is a link is not walked)
                if rel in DIRS and st is not None and stat.S_ISDIR(st.st_mode):
                    try:
                        top = os.open(name, DIR_FLAGS, dir_fd=fd)
                    except OSError:
                        continue
                    try:
                        self.walk(top, rel, 1)
                    finally:
                        os.close(top)
            finally:
                os.close(fd)
        os.close(self.home)
        return self.out


def manifest(home, copies=None):
    return Manifest(home, copies).lines()


def path_of(line):
    parts = line.split(" ", 4)
    return parts[4] if len(parts) == 5 else None


def changed(old, new):
    """The paths whose line differs between two manifests (changed, new or gone)"""
    return {p for p in (path_of(line) for line in set(old) ^ set(new)) if p}


# ---- users

def home_of(user):
    # (a test gives its own pretend home; never with root's real record)
    if TEST and os.environ.get("WREN_WATCH_HOME"):
        return os.environ["WREN_WATCH_HOME"]
    try:
        return pwd.getpwnam(user).pw_dir
    except KeyError:
        return ""


def user_ok(user):
    return bool(re.fullmatch(r"[a-z_][a-z0-9_-]{0,31}", user)) and bool(home_of(user))


def as_user(user, cmd, **kw):
    """CMD with USER's rights, not root's. A program running as USER can plant a link in those
    files, to a root-only file or the signing key: as USER, it reads only what USER may read
    anyway"""
    if os.getuid() == 0:
        pw = pwd.getpwnam(user)
        cmd = ["setpriv", "--reuid=%d" % pw.pw_uid, "--regid=%d" % pw.pw_gid, "--clear-groups", "--no-new-privs"] + cmd
    return subprocess.run(cmd, **kw)


def me(cmd):
    if os.getuid() == 0:
        user = os.environ.get("SUDO_USER", "")
        if not user or user == "root":
            die("run it through sudo, as yourself: sudo wren-watch " + cmd)
        return user
    return pwd.getpwuid(os.getuid()).pw_name


# ---- check, accept

def check(user):
    """Compare, and say what changed in RUN/USER (readable by USER only)"""
    if not user_ok(user):
        return
    record = os.path.join(STATE, user, "manifest")
    if not os.path.isfile(record):
        return
    found = changed(read_lines(record), manifest(home_of(user)))
    os.makedirs(RUN, mode=0o755, exist_ok=True)
    os.chmod(RUN, 0o755)
    # Once listed, a file stays listed until you accept it, even if it is back as it was by the
    # next look: a program could change it, wait for it to be used, and put it back. The list
    # is kept in root's record too (STATE/USER/listed), which a reboot doesn't empty as it
    # does RUN; RUN has the copy the bar reads.
    listed = os.path.join(STATE, user, "listed")
    bar = os.path.join(RUN, user)
    found |= set(read_lines(listed)) | set(read_lines(bar))
    found.discard("")
    if found:
        text = "".join(p + "\n" for p in sorted(found))
        write_file(listed, text, 0o600)
        write_file(bar, text, 0o640)
        if not TEST:
            os.chown(bar, -1, pwd.getpwnam(user).pw_gid)
    else:
        remove(bar)
        remove(listed)


def accept(user, paths):
    """The record as the files are now (all, or those paths)"""
    home = home_of(user)
    mine = os.path.join(STATE, user)
    for d in (STATE, mine):
        os.makedirs(d, mode=0o700, exist_ok=True)
        os.chmod(d, 0o700)
    copies = os.path.join(mine, "copies")
    record = os.path.join(mine, "manifest")
    if not paths:
        shutil.rmtree(copies + ".new", ignore_errors=True)
        now = manifest(home, copies + ".new")
        shutil.rmtree(copies, ignore_errors=True)
        if os.path.isdir(copies + ".new"):
            os.rename(copies + ".new", copies)
        write_file(record, "".join(line + "\n" for line in now), 0o600)
    else:
        # only those paths take their new line; the rest keep the old one
        part = copies + ".part"
        shutil.rmtree(part, ignore_errors=True)
        now = manifest(home, part)
        # the copies too: only those paths' (for diff, the others stay as they were)
        for p in paths:
            remove(os.path.join(copies, p))
            src = os.path.join(part, p)
            if os.path.isfile(src):
                dest = os.path.join(copies, p)
                os.makedirs(os.path.dirname(dest), mode=0o700, exist_ok=True)
                shutil.copyfile(src, dest)
                os.chmod(dest, 0o600)
        shutil.rmtree(part, ignore_errors=True)
        old = {path_of(line): line for line in read_lines(record) if path_of(line)}
        new = {path_of(line): line for line in now if path_of(line)}
        want = set(paths)
        kept = []
        for p in sorted(set(old) | set(new)):
            src = new if p in want else old
            if p in src:
                kept.append(src[p] + "\n")
        write_file(record, "".join(kept), 0o600)
    # what was accepted leaves the list, both copies (check keeps the rest listed)
    for f in (os.path.join(RUN, user), os.path.join(mine, "listed")):
        if not paths:
            remove(f)
        elif os.path.isfile(f):
            left = [line for line in read_lines(f) if line not in set(paths)]
            if left:
                write_file(f, "".join(line + "\n" for line in left), os.stat(f).st_mode & 0o777)
            else:
                remove(f)
    check(user)


# ---- the commands

def cmd_status():
    user = me("status")
    listed = read_lines(os.path.join(RUN, user))
    if not listed:
        say("Nothing changed in your startup files since you last accepted.")
        return
    say("Changed since you last accepted (if it wasn't you or Wren Linux, something running as you did it):")
    for rel in listed:
        say("  ~/" + rel)
    say("See the changes: wren-watch diff    Take them as yours: wren-watch accept    (sudo asks)")
    say("If it was neither you nor Wren Linux: don't accept it, don't type passwords, and take this")
    say("computer off the network. docs/opsec.md, \"When something goes wrong\", says what next.")


def sudo_self(*args):
    """Through sudo, as root's copy by its path: sudo asks for your password"""
    try:
        os.execvp(SUDO, [SUDO, "--", SELF] + list(args))
    except OSError as e:
        die("couldn't run %s: %s" % (SUDO, e.strerror))


def cmd_diff():
    # it needs root's copies: sudo asks for your password
    if os.getuid() != 0 and not TEST:
        sudo_self("diff")
    user = me("diff")
    home = home_of(user)
    listed = read_lines(os.path.join(RUN, user))
    if not listed:
        say("Nothing changed.")
        return
    quiet = {"stdout": subprocess.DEVNULL, "stderr": subprocess.DEVNULL}
    for rel in listed:
        live = os.path.join(home, rel)
        copy = os.path.join(STATE, user, "copies", rel)
        before, now = "--label=before: ~/" + rel, "--label=now: ~/" + rel
        # (the record's copy is root's; the file as it is now is read as you)
        if as_user(user, ["test", "-L", live], **quiet).returncode == 0:
            target = as_user(user, ["readlink", "--", live], stdout=subprocess.PIPE, stderr=subprocess.DEVNULL).stdout
            say("~/%s is a link to %s (not followed)" % (rel, shlex.quote(os.fsdecode(target.rstrip(b"\n")))))
        elif (as_user(user, ["test", "-f", live], **quiet).returncode == 0 and
              as_user(user, ["test", "-r", live], **quiet).returncode == 0):
            content = as_user(user, ["cat", "--", live], stdout=subprocess.PIPE, stderr=subprocess.DEVNULL).stdout
            d = subprocess.run(["diff", "-u", "-N", before, now, copy, "-"], input=content,
                               stdout=subprocess.PIPE, stderr=subprocess.DEVNULL).stdout
            if d:
                sys.stdout.write(shown(d.decode("utf-8", "surrogateescape")))
            else:
                say("~/%s changed since you last accepted, and is back as it was then (or only its permissions changed)" % rel)
        elif os.path.isfile(copy):
            d = subprocess.run(["diff", "-u", before, "--label=now: ~/%s (gone)" % rel, copy, "/dev/null"],
                               stdout=subprocess.PIPE, stderr=subprocess.DEVNULL).stdout
            sys.stdout.write(shown(d.decode("utf-8", "surrogateescape")))
        else:
            say("~/%s (not a text file, or gone)" % rel)
        sys.stdout.flush()


def cmd_accept(args):
    # it changes root's record: sudo asks for your password
    if os.getuid() != 0 and not TEST:
        sudo_self("accept", *args)
    user = me("accept")
    home = home_of(user)
    paths = []
    for a in args:
        a = a.removeprefix(home + "/").removeprefix("~/")
        if a.startswith("/") or ".." in a:
            die("\"%s\": a file in your home, please (as wren-watch lists it)." % shown(a))
        paths.append(a)
    accept(user, paths)
    say("Accepted. The record is your startup files as they are now%s." % (" (%s)" % " ".join(paths) if paths else ""))


def cmd_check_all():
    if os.getuid() != 0 and not TEST:
        die("that is root's timer's.")
    failed = False
    users = sorted(d for d in os.listdir(STATE) if os.path.isdir(os.path.join(STATE, d))) if os.path.isdir(STATE) else []
    # one user's trouble doesn't keep the others from being looked at; the timer says it failed
    for user in users:
        try:
            check(user)
        except Exception as e:
            print("wren-watch: checking %s: %s" % (shown(user), shown(str(e))), file=sys.stderr)
            failed = True
    if failed:
        sys.exit(1)


def cmd_help():
    with open(__file__, encoding="utf-8") as f:
        next(f)
        for line in f:
            if not line.startswith("#"):
                break
            sys.stdout.write(re.sub(r"^# ?", "", line))


def main(argv):
    # (what it prints is plain text in any language setting: names that aren't UTF-8 are
    # written escaped, never as an error that would end the look)
    sys.stdout.reconfigure(encoding="utf-8", errors="surrogateescape")
    cmd, args = (argv[0], argv[1:]) if argv else ("", [])
    if cmd in ("", "status"):
        cmd_status()
    elif cmd == "diff":
        cmd_diff()
    elif cmd == "accept":
        cmd_accept(args)
    elif cmd == "manifest":
        for line in manifest(os.environ.get("HOME", "")):
            sys.stdout.write(line + "\n")
    elif cmd == "check-all":
        cmd_check_all()
    elif cmd in ("-h", "--help"):
        cmd_help()
    else:
        die("usage: wren-watch [status | diff | accept [FILE...] | manifest | check-all]")


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