#!/usr/bin/env python3
"""
omv-imagecheck - monitor-only Docker image update checker for OpenMediaVault.

For every RUNNING container it compares the image's local registry digest
(RepoDigest, i.e. what was actually pulled) against the current digest the
registry advertises for that same tag. If they differ, a newer image is
available. It NEVER pulls, recreates, or changes anything - it only reports.

The result is written atomically to a JSON cache that the ImageCheck RPC
service reads. Run daily from an OMV scheduled task, or on demand via
ImageCheck.refresh.

Usage:
    omv-imagecheck --refresh     # do the check and update the cache
    omv-imagecheck --print       # print the current cache (no check)

Exit code is 0 on a completed check even if some images could not be looked
up (those are recorded as per-image errors, not a global failure).
"""

import argparse
import json
import os
import subprocess
import sys
import tempfile
import time

CACHE_FILE = "/var/lib/openmediavault/imagecheck.json"

__version__ = "8.1.1"

# Per-command timeouts (seconds). Registry lookups can be slow.
DOCKER_TIMEOUT = 30
REGISTRY_TIMEOUT = 45

# regctl (regclient) — installed by install.sh into /usr/local/bin. Queries
# registries directly (no pull, no Docker daemon) for the current tag digest.
REGCTL = os.environ.get("OMV_IMAGECHECK_REGCTL", "regctl")


def run(args, timeout):
    """Run a command, return (returncode, stdout, stderr) as stripped text."""
    try:
        p = subprocess.run(
            args,
            capture_output=True,
            text=True,
            timeout=timeout,
        )
        return p.returncode, p.stdout.strip(), p.stderr.strip()
    except FileNotFoundError:
        return 127, "", "%s: command not found" % (args[0],)
    except subprocess.TimeoutExpired:
        return 124, "", "timed out after %ss" % (timeout,)
    except Exception as e:  # pragma: no cover - defensive
        return 1, "", str(e)


def docker_available():
    rc, _out, _err = run(["docker", "version", "--format", "{{.Server.Version}}"],
                          DOCKER_TIMEOUT)
    return rc == 0


def regctl_available():
    rc, _out, _err = run([REGCTL, "version"], DOCKER_TIMEOUT)
    return rc == 0


def running_containers():
    """Return list of dicts: {name, image} for each running container."""
    rc, out, err = run(["docker", "ps", "--format", "{{json .}}"], DOCKER_TIMEOUT)
    if rc != 0:
        return None, (err or "docker ps failed")
    items = []
    for line in out.splitlines():
        line = line.strip()
        if not line:
            continue
        try:
            obj = json.loads(line)
        except ValueError:
            continue
        name = obj.get("Names") or obj.get("Name") or ""
        image = obj.get("Image") or ""
        items.append({"name": name, "image": image})
    return items, None


def normalize_ref(image):
    """
    Turn a `docker ps` image field into a ref we can inspect remotely.
    Returns None for refs we cannot meaningfully check (pinned by digest,
    bare image IDs, dangling images).
    """
    if not image or image == "<none>:<none>" or image.startswith("<none>"):
        return None
    # Already pinned to a digest - there is nothing newer to move to.
    if "@sha256:" in image:
        return None
    # A bare 12/64-hex image ID (no tag, no slash) is not a registry ref.
    bare = image.split(":", 1)[0]
    if "/" not in image and ":" not in image and _looks_like_id(bare):
        return None
    return image


def _looks_like_id(s):
    if len(s) not in (12, 64):
        return False
    try:
        int(s, 16)
        return True
    except ValueError:
        return False


def local_digest(image_ref, repo):
    """
    The registry digest of the locally-present image (what was pulled), read
    from RepoDigests. Uses `docker image inspect` with a JSON format string so
    there are no shell/Go-template escaping pitfalls, and picks the RepoDigest
    whose repository matches image_ref (an image can carry digests for several
    repos). Forces image lookup, never a same-named container.
    Returns 'sha256:...' or None.
    """
    rc, out, err = run(
        ["docker", "image", "inspect", "--format", "{{json .RepoDigests}}",
         image_ref],
        DOCKER_TIMEOUT,
    )
    if rc != 0 or not out:
        return None
    try:
        digests = json.loads(out)
    except ValueError:
        return None
    if not isinstance(digests, list) or not digests:
        return None
    # Prefer the entry whose repo matches this image's repo.
    for entry in digests:
        if isinstance(entry, str) and "@sha256:" in entry:
            entry_repo, entry_digest = entry.split("@", 1)
            if entry_repo == repo:
                return entry_digest
    # Fall back to the first digest present.
    first = digests[0]
    if isinstance(first, str) and "@" in first:
        return first.split("@", 1)[1]
    return None


def remote_digest(image_ref):
    """
    The digest the registry currently advertises for image_ref's tag, fetched
    via regctl (regclient) WITHOUT pulling the image and WITHOUT needing the
    Docker daemon. `regctl image digest` returns the index/manifest digest the
    tag points to - what RepoDigests holds too - so the two are directly
    comparable. Private registries use the host's ~/.docker/config.json creds.
    Returns (digest, None) or (None, error_message).
    """
    rc, out, err = run(
        [REGCTL, "image", "digest", image_ref],
        REGISTRY_TIMEOUT,
    )
    if rc == 0 and out.startswith("sha256:"):
        return out.strip(), None
    msg = err or out or "registry lookup failed"
    # Keep the message compact for the UI.
    msg = msg.splitlines()[0][:200] if msg else "registry lookup failed"
    return None, msg


def split_tag(image_ref):
    """Best-effort (repo, tag) split for display. Defaults tag to 'latest'."""
    # Strip digest if somehow present.
    ref = image_ref.split("@", 1)[0]
    # A ':' after the last '/' is the tag; a ':' in the host part is a port.
    last_slash = ref.rfind("/")
    tail = ref[last_slash + 1:]
    if ":" in tail:
        repo_tail, tag = tail.rsplit(":", 1)
        repo = ref[:last_slash + 1] + repo_tail
        return repo, tag
    return ref, "latest"


def do_check():
    started = int(time.time())
    result = {
        "generatedAt": started,
        "checked": 0,
        "updateCount": 0,
        "error": "",
        "images": [],
    }

    if not docker_available():
        result["error"] = ("Docker is not available to this user. The checker "
                            "must run as root (or a docker-group user).")
        print(result["error"])
        return result

    if not regctl_available():
        result["error"] = ("regctl (regclient) is not installed. Re-run the "
                            "plugin installer, which downloads it to /usr/local/bin.")
        print(result["error"])
        return result

    containers, err = running_containers()
    if err:
        result["error"] = err
        print(err)
        return result

    print("Checking %d running container(s) for image updates..."
          % len(containers))

    # Cache remote/local digests per unique image ref so we hit the registry
    # only once per image no matter how many containers share it.
    remote_cache = {}
    local_cache = {}
    rows = []
    update_count = 0
    checked = 0

    for c in containers:
        image = c["image"]
        ref = normalize_ref(image)
        repo, tag = split_tag(image)

        if ref is None:
            rows.append({
                "container": c["name"],
                "image": image,
                "repo": repo,
                "tag": tag,
                "local": "",
                "remote": "",
                "updateAvailable": False,
                "error": "Not a checkable registry tag (local build or digest-pinned)",
            })
            print("  %s (%s): skipped - not a checkable tag" % (c["name"], image))
            continue

        checked += 1

        if ref not in local_cache:
            local_cache[ref] = local_digest(ref, repo)
        ld = local_cache[ref]

        if ref not in remote_cache:
            remote_cache[ref] = remote_digest(ref)
        rd, rerr = remote_cache[ref]

        row = {
            "container": c["name"],
            "image": image,
            "repo": repo,
            "tag": tag,
            "local": ld or "",
            "remote": rd or "",
            "updateAvailable": False,
            "error": "",
        }

        if rerr:
            row["error"] = rerr
            print("  %s (%s): lookup error - %s" % (c["name"], ref, rerr))
        elif not ld:
            row["error"] = "No local registry digest (image not pulled from a registry?)"
            print("  %s (%s): no local digest" % (c["name"], ref))
        elif rd and ld != rd:
            row["updateAvailable"] = True
            update_count += 1
            print("  %s (%s): UPDATE AVAILABLE" % (c["name"], ref))
        else:
            print("  %s (%s): up to date" % (c["name"], ref))

        rows.append(row)

    result["checked"] = checked
    result["updateCount"] = update_count
    result["images"] = rows
    print("Done. %d update(s) available across %d checked container(s)."
          % (update_count, checked))
    return result


def write_cache(data):
    directory = os.path.dirname(CACHE_FILE)
    try:
        os.makedirs(directory, exist_ok=True)
    except OSError:
        pass
    # Atomic write: temp file in the same dir, then rename.
    fd, tmp = tempfile.mkstemp(dir=directory, prefix=".imagecheck.", suffix=".tmp")
    try:
        with os.fdopen(fd, "w") as f:
            json.dump(data, f, separators=(",", ":"))
            f.flush()
            os.fsync(f.fileno())
        os.replace(tmp, CACHE_FILE)
        try:
            os.chmod(CACHE_FILE, 0o644)
        except OSError:
            pass
    finally:
        if os.path.exists(tmp):
            try:
                os.remove(tmp)
            except OSError:
                pass


def read_cache():
    try:
        with open(CACHE_FILE) as f:
            return json.load(f)
    except (OSError, ValueError):
        return {
            "generatedAt": 0, "checked": 0, "updateCount": 0,
            "error": "", "images": [],
        }


def main():
    parser = argparse.ArgumentParser(description="OMV Docker image update checker")
    parser.add_argument("--version", action="version",
                        version="omv-imagecheck " + __version__)
    group = parser.add_mutually_exclusive_group()
    group.add_argument("--refresh", action="store_true",
                       help="run the check and update the cache")
    group.add_argument("--print", dest="do_print", action="store_true",
                       help="print the current cache without checking")
    args = parser.parse_args()

    if args.do_print:
        print(json.dumps(read_cache(), indent=2))
        return 0

    # Default action is --refresh.
    data = do_check()
    write_cache(data)
    return 0


if __name__ == "__main__":
    sys.exit(main())
