ctre

#!/usr/bin/env python3
# ctre (part of ossobv/vcutil) // wdoekes/2023-2024 // Public Domain
#
# Enhanced CTR (containerd CLI)
#
# Right now it only does pruning of unused images. This is useful in
# Kubernetes (K8S) environments where stale content can pile up.
#
# Usage:
#
#     ctre prune                # lists all images to prune
#     ctre prune -f             # prunes the images
#     ctre record-image-starts  # daemon process that keeps track of starts
#
# This differs from 'crictl prune' because it records image use first.
# Making sure that we're not removing images that were recently used.
# Recording is done in /var/spool/ctre/image_starts by this application
# run in daemon mode.
#
# BEWARE: stoptime can record the stopping of the _previous_ container after
# starting the next one. This can confuse you into thinking a container is
# started/stopped quickly.
#
from collections import defaultdict, namedtuple
from datetime import datetime
from json import loads
from re import compile
from os import makedirs, rename
from os.path import dirname, basename
from subprocess import PIPE, STDOUT, CalledProcessError, Popen, check_output
from sys import argv, stderr
from time import sleep, time
import warnings
from warnings import warn


KEEP_IF_YOUNGER_THAN = (14 * 86400)
NULL_TIME = '----------'
TABSPLIT_RE = compile(r'\t+')
SIZE_UNITS = {
    'B': 1,
    'kB': 1024,
    'MB': 1024 * 1024,
    'GB': 1024 * 1024 * 1024,
    ' KiB': 1024,
    ' MiB': 1024 * 1024,
    ' GiB': 1024 * 1024 * 1024,
}

IMAGE_STARTS_DB = '/var/spool/ctre/image_starts'

Blob = namedtuple('Blob', 'ns digest size age')
Event = namedtuple('Event', 'unixtime ns action data')
Image = namedtuple('Image', 'ns digest ref size age')


class Skip(ValueError):
    pass


def split_by_header(lines):
    """
    Return lines (with header) as dictionaries

    In ["KEY1  KEY2  KEY3", "VAL 1 VAL 2 VAL 3"]
    Out [{"KEY1": "VAL 1", "KEY2": "VAL 2", "KEY3": "VAL 3"}]

    .. or ..

    In ["KEY1<TAB>KEY2<TAB>KEY3", "VAL 1<TAB>VAL 2<TAB>VAL 3"]
    Out [{"KEY1": "VAL 1", "KEY2": "VAL 2", "KEY3": "VAL 3"}]
    """
    def as_value(value):
        if value == '-':
            return None
        return value

    header = lines[0]
    parts = header.strip().split()
    ret = []

    if '\t' in header:
        for line in lines[1:]:
            data = TABSPLIT_RE.split(line)
            assert len(data) == len(parts), (parts, data)
            item = {}
            for n, part in enumerate(parts):
                item[part] = as_value(data[n])
            ret.append(item)

    else:
        sizes = []
        last_idx = 0
        for next_part in parts[1:]:
            idx = header.find(next_part, last_idx)
            assert idx > 0, (next_part, idx, header)
            assert (
                len(header[idx:]) == len(next_part) or
                header[idx + len(next_part)] in ' \t'), (
                    next_part, idx, header)
            assert header[idx - 1] in ' \t', (next_part, idx, header)
            sizes.append((last_idx, idx))
            last_idx = idx
        sizes.append((last_idx, -1))

        for line in lines[1:]:
            item = {}
            for part, size in zip(parts, sizes):
                item[part] = as_value(line[size[0]:size[1]].rstrip())
            ret.append(item)

    return ret


def from_size(s):
    size = []
    for n, ch in enumerate(s):
        if ch not in '0123456789.':
            size = float(''.join(size))
            unit = s[n:]
            break
        else:
            size.append(ch)

    return int(SIZE_UNITS[unit] * size)


def from_age(s):
    n, tail = s.split(' ', 1)

    if n.isdigit():
        n = int(n)
        if tail == 'years':
            return int(n * 86400 * 365.25)
        if tail == 'months':
            return int(n * 86400 * 30.4375)
        if tail == 'weeks':
            return (n * 86400 * 7)
        if tail == 'days':
            return (n * 86400)
        if tail == 'hours':
            return (n * 3600)
        if tail == 'minutes':
            return (n * 60)
        if tail == 'seconds':
            return n

    if n == 'About':
        if tail == 'an hour':
            return 3600
        if tail == 'a minute':
            return 60

    raise NotImplementedError(s)


def namespaces():
    out = check_output(['ctr', 'namespace', 'ls', '--quiet']).decode().strip()
    return out.split('\n')


def blobs(ns):
    def dict_to_blob(ns, d):
        assert len(d) == 4, d
        return Blob(
            ns=ns,
            digest=d['DIGEST'],
            size=from_size(d['SIZE']),
            age=from_age(d['AGE']),
        )

    out = check_output(['ctr', '-n', ns, 'content', 'ls']).decode().strip()
    # DIGEST  SIZE  AGE  LABELS
    # sha..   123K  1 mo ..
    items = split_by_header(out.split('\n'))
    blobs = [dict_to_blob(ns, item) for item in items]
    return dict((blob.digest, blob) for blob in blobs)


def images(ns):
    def dict_to_image(ns, item, blob_dict):
        blob = blob_dict[item['DIGEST']]
        # assert from_size(item['SIZE']) == blob.size, (item['SIZE'], blob)
        return Image(
            ns=ns,
            digest=blob.digest,
            ref=item['REF'],
            size=from_size(item['SIZE']),
            age=blob.age,
        )

    blob_dict = blobs(ns)

    out = check_output(['ctr', '-n', ns, 'image', 'ls']).decode().strip()
    # REF  TYPE  DIGEST  SIZE  PLATFORMS  LABELS
    # name       sha..   1 MiB ...        ...
    items = split_by_header(out.split('\n'))
    images = [dict_to_image(ns, item, blob_dict) for item in items]
    return images


def images_not_in_use(ns):
    out = check_output(['ctr', '-n', ns, 'container', 'ls']).decode().strip()
    # CONTAINER  IMAGE  RUNTIME
    # can contain empty images "-" when they are not linked through containerd
    # (as seen with moby/dockerd)
    items = split_by_header(out.split('\n'))
    images_in_use = set([item['IMAGE'] for item in items])

    image_list = images(ns)
    image_dict = dict((image.ref, image.digest) for image in image_list)
    image_digest_in_use = set()
    for image_ref in images_in_use:
        if image_ref is not None:
            try:
                image_digest_in_use.add(image_dict[image_ref])
            except KeyError:
                warn(f'Container ref {image_ref!r} not found in images')

    # Order by digest
    by_digest = defaultdict(list)
    for image in image_list:
        by_digest[image.digest].append(image)

    # Prepare results
    not_in_use = []
    for image_digest, images_for_digest in by_digest.items():
        assert images_for_digest, images_for_digest

        if image_digest in image_digest_in_use:
            # For those that _are_ in use, report about those with multiple
            # labels. We expect:
            # - img:tag
            # - img@sha256:abc..
            # - sha256:abc..
            img_tag = []
            img_sha = []
            for image in images_for_digest:
                if image.ref.startswith('sha256:'):
                    # assert image.ref == image_digest, (ref, image_digest)
                    assert len(image.ref) == 71, image.ref
                elif '@' in image.ref:
                    baseref, digest = image.ref.split('@')
                    assert digest == image_digest, (baseref, image_digest)
                    assert ':' not in baseref, image.ref
                    img_sha.append(baseref)
                elif ':' in image.ref:
                    baseref, tag = image.ref.split(':', 1)
                    img_tag.append(baseref)
                else:
                    assert False, image.ref

            # Do not report duplicate tags. But do report duplicate names.
            if len(set(img_sha) | set(img_tag)) > 1:
                all_refs = ' + '.join(sorted(
                    image.ref for image in images_for_digest))
                image = images_for_digest[0]
                warn(
                    f'Dupe refs found for ns {image.ns} image {image.digest}, '
                    f'keeping all: {all_refs}')
        else:
            # For those _not_ in use, add to our response.
            not_in_use.extend(images_for_digest)

    return not_in_use


def prunable_images(ns):
    ret = []
    for image in images_not_in_use(ns):
        # Check image creation time.
        if image.age >= KEEP_IF_YOUNGER_THAN:
            ret.append(image)
    return ret


def get_all_prunable_images():
    images = []
    for ns in namespaces():
        for image in prunable_images(ns):
            images.append((ns, image.ref))
    return images


def prunable_images_after_double_check():
    prunable = set(get_all_prunable_images())
    tnow = time()
    tmax = tnow + 15

    # Load data collected by ctre record-image-starts daemon
    kv = load_image_starts(max_age=KEEP_IF_YOUNGER_THAN)
    used_to_start_a_container_recently = set(kv.keys())
    prunable -= used_to_start_a_container_recently

    while tnow <= tmax and prunable:
        left = int(tmax - tnow)

        if stderr.isatty():
            msg = 'double check to avoid cleaning mid-restart'
            stderr.write(f'\x1b[1K\r({msg}... {left} secs)')
            stderr.flush()

        sleep(1.5)
        tnow = time()

        prunable_next = set(get_all_prunable_images())
        prunable &= prunable_next  # only keep those in both lists

    if stderr.isatty():
        stderr.write('\x1b[1K\r')
        stderr.flush()

    return list(sorted(prunable))


def remove_images(images):
    for ns, ref in images:
        ret = check_output(
            ['ctr', '-n', ns, 'image', 'rm', ref]).decode().strip()
        print('(deleted)', ns, ret)


def remove_containers():
    # Simply remove all containers. Every container that fails is still
    # running. Don't worry about that.
    for ns in namespaces():
        containers = check_output(
            ['ctr', '-n', ns, 'containers', 'list', '--quiet']
            ).decode().strip()
        if not containers:
            continue

        container_objects = []
        for container_id in containers.split('\n'):
            try:
                container_objects.append(_get_container_info(ns, container_id))
            except Skip:
                pass
            except (AttributeError, KeyError) as e:
                print('(container info error)', ns, container_id, e)
                breakpoint()

        container_objects.sort(key=(
            lambda x: (x['k8s.ns'], x['k8s.pod'], x['updated'])))
        for obj in container_objects:
            try:
                out = check_output(
                    ['ctr', '-n', ns, 'containers', 'remove', obj['id']],
                    stderr=STDOUT,
                    ).decode()
            except CalledProcessError:
                # TODO: Hide this.
                print(
                    '(keeping container)', ns, obj['k8s.ns'], obj['k8s.pod'],
                    obj['k8s.container'], obj['created'])
            else:
                # TODO: Maybe not list the timestamps..
                print(
                    '(deleted container)', ns, obj['k8s.ns'], obj['k8s.pod'],
                    obj['k8s.container'], obj['created'])


def _get_container_info(ns, container_id):
    out = check_output(
        ['ctr', '-n', ns, 'containers', 'info', container_id]
        ).decode().strip()
    json = loads(out)
    assert json['ID'] == container_id

    if 'Labels' not in json or json['Labels'] is None:
        assert json['Image'] == '', (ns, container_id, json['Image'])
        print('(no labels in container, skipping)', ns, container_id)
        raise Skip

    return {
        'id': container_id,
        'k8s.ns': json['Labels'].get(
            'io.kubernetes.pod.namespace', container_id),
        'k8s.pod': json['Labels'].get(
            'io.kubernetes.pod.name', container_id),
        'k8s.container': json['Labels'].get(
            'io.kubernetes.container.name', '[{}]'.format(
                json['Labels'].get('io.cri-containerd.kind'))),
        # TODO: Use timestamps?
        'created': json['CreatedAt'],
        'updated': json['UpdatedAt'],
    }


def prune_content_references():
    # Not entirely sure what this does, but it shrinks the content down.
    # We'll likely want to do this last.
    for ns in namespaces():
        before = check_output(
            ['ctr', '-n', ns, 'content', 'list', '--quiet']
            ).decode().count('\n')
        ret = check_output(
            ['ctr', '-n', ns, 'content', 'prune', 'references']
            ).decode().strip()
        if ret != '':
            print('(unexpected ctr content prune references output)', ns, ret)
        after = check_output(
            ['ctr', '-n', ns, 'content', 'list', '--quiet']
            ).decode().count('\n')
        print('(pruned content)', ns, (before - after))


def container_image_ref(ns, container):
    out = (
        check_output(['ctr', '-n', ns, 'containers', 'info', container])
        .decode('utf-8').strip())
    data = loads(out)
    return data['Image']


def readlines(fp):
    ret = []
    while True:
        ch = fp.read(1)  # or os.read(fp.fileno(), 1)
        if not ch:
            raise ValueError('got EOF from ctr event')
        ret.append(ch)
        if ch == b'\n':
            yield b''.join(ret)
            ret = []


def event_lines(it):
    for line in it:
        line = line.decode('utf-8', 'replace')
        date, time, toff, tzone, ns, action, data = line.split(' ', 6)
        time = time[0:15]  # strip to microseconds: "%H:%M:%S.%fffff"
        when = datetime.strptime(
            f'{date} {time} {toff} {tzone}', '%Y-%m-%d %H:%M:%S.%f %z %Z')
        unixtime = when.timestamp()
        data = loads(data)
        yield Event(unixtime=unixtime, ns=ns, action=action, data=data)


def from_kv_old(s):
    """
    <starttime> <namespace> <image>  # old
    <starttime> <image>              # oldest
    """
    ret = {}
    for line in s.split('\n'):
        v, k = line.split(' ', 1)
        if ' ' in k:
            k = tuple(k.split())
        ret[k] = v

    return ret


def from_kv(s):
    """
    <starttime> <stoptime|-> <namespace> <image>

    BEWARE: stoptime can record the stopping of the _previous_ container after
    starting the next one. This can confuse you into thinking a container is
    started/stopped quickly.
    """
    s = s.rstrip()
    if not s:
        return {}

    first_line = s.split('\n', 1)[0]
    first_line_items = first_line.split(' ', 3)
    if len(first_line_items) < 4 or not (
            first_line_items[1].isdigit() or
            first_line_items[1] == NULL_TIME):
        oldret = from_kv_old(s)
        ret = dict((k, (int(v), None)) for k, v in oldret.items())
        return ret

    ret = {}
    for line in s.split('\n'):
        start, stop, namespace, image = line.split(' ', 3)
        k = (namespace, image)  # ctr namespace
        v = (int(start), None if stop == NULL_TIME else int(stop))
        ret[k] = v

    return ret


def to_kv(kv):
    ret = []
    for k, v in sorted(kv.items()):
        flatk = ' '.join(k)
        starttime, stoptime = v
        if stoptime is None:
            stoptime = NULL_TIME
        ret.append(f'{starttime} {stoptime} {flatk}\n')
    return ''.join(ret)


def load_image_starts(max_age):
    filename = IMAGE_STARTS_DB

    # Load
    try:
        with open(filename, 'r') as fp:
            all_ = fp.read()
    except FileNotFoundError:
        return {}

    kv = from_kv(all_)

    # Prune
    old = time() - max_age
    to_rm = set()
    for image_ref, (start, stop) in kv.items():
        # Check start/stop time of containers with this image.
        if start < old and (stop is None or stop < old):
            to_rm.add(image_ref)
    for rm in to_rm:
        del kv[rm]

    return kv


def update_image_startstop(starttime, stoptime, ns, image_ref):
    filename = IMAGE_STARTS_DB
    temp = f'{filename}.tmp'

    # Load
    kv = load_image_starts(
        max_age=max(KEEP_IF_YOUNGER_THAN, 14 * 86400))  # prune to 14 days

    # Add
    try:
        prev_starttime, prev_stoptime = kv[(ns, image_ref)]
    except KeyError:
        prev_starttime = prev_stoptime = None
    starttime = starttime or prev_starttime or stoptime
    stoptime = stoptime or prev_stoptime
    kv[(ns, image_ref)] = (starttime, stoptime)

    # Save
    all_ = to_kv(kv)
    try:
        with open(temp, 'w') as fp:
            fp.write(all_)
    except FileNotFoundError:
        makedirs(dirname(temp))
        with open(temp, 'w') as fp:
            fp.write(all_)
    rename(temp, filename)


def record_image_starts():
    print('Listening for container start events')
    with Popen(['ctr', 'event'], stdout=PIPE) as fp:
        ignored_events = set()

        for event in event_lines(readlines(fp.stdout)):
            try:
                handle_event(event, ignored_events, debug=False)
            except Exception as e:
                print(f'ERROR event: {event!r} -- {e}', file=stderr)
                raise


def handle_event(event, ignored_events, debug=False):
    # Auto-debug stuff during startup: print all.
    if event.action not in ignored_events:
        # CREATE POD
        # - /snapshot/prepare   Create snapshot [key]
        # - /containers/create  Create container (for pod)  [id]
        # - /containers/update  Set container info/config   [id]
        # - /tasks/create       Create task (start pod) [container_id]
        # - /tasks/start        Start task (start pod)  [container_id]
        # END POD
        # - /tasks/exit         Ends container (for pod)
        #                       [id and container_id (same value)]
        # - /snapshot/remove    Clear snapshot [key]
        # - /containers/delete  Clear container [id]
        #
        # CREATE EXEC
        # - /tasks/exec-added   Create exec [container_id, exec_id]
        # - /tasks/exec-started Start exec [container_id, exec_id]
        # END EXEC
        # - /tasks/exit         Ends container (for exec)
        #                       [id and container_id (different value)]
        #
        # OTHER
        # - /images/update      Update image info after pod start
        #                       (add 'io.cri-containerd.image' label)
        print(
            f'D (first visual of event action {event.action!r}) '
            f'data: {event.data!r}')
        ignored_events.add(event.action)

    # In container.io env, we first see a container create event.
    if event.action == '/containers/create':
        unixtime = int(event.unixtime)
        if 'image' in event.data:
            container_id = event.data['id']
            image_ref = event.data['image']
            update_image_startstop(unixtime, None, event.ns, image_ref)
            print('C', unixtime, event.ns, image_ref, (
                event.data if debug else ''))
        else:
            print('C', unixtime, event.ns, 'NO_IMAGE_REF', (
                event.data if (event.ns != 'moby' or debug) else ''))

    # Tasks create/start is how we see PODs getting started, both in
    # docker and pure container.io environments.
    elif event.action in (
            # Typical POD startup on ZFS:
            # /snapshot/prepare {"key":"9bd","parent":"sha256:dee"}
            # /containers/create {"id":"9bd","image":"k8s.gcr.io/...
            # /containers/update {"id":"9bd","image":"k8s.gcr.io/...
            # /tasks/create {"container_id":"9bd","bundle":"/run/...
            #   "rootfs":[{"type":"zfs",
            #   "source":"local-storage/containerd/43307"}],...
            # /tasks/start {"container_id":"9bd","pid":2286326}
            # /images/update {'name': 'docker.io/..', 'labels':..
            '/tasks/create'):
        container_id = event.data['container_id']
        unixtime = int(event.unixtime)
        image_ref = container_image_ref(event.ns, container_id)
        if image_ref:
            update_image_startstop(unixtime, None, event.ns, image_ref)
            print('T', unixtime, event.ns, image_ref, (
                event.data if debug else ''))
        else:
            print('T', unixtime, event.ns, 'NO_IMAGE_REF', (
                event.data if (event.ns != 'moby' or debug) else ''))

    # Stop/end tasks
    elif event.action == '/tasks/exit' and (
            event.data['id'] == event.data['container_id']):
        container_id = event.data['container_id']
        unixtime = int(event.unixtime)
        image_ref = container_image_ref(event.ns, container_id)
        if image_ref:
            update_image_startstop(None, unixtime, event.ns, image_ref)
            print('t', unixtime, event.ns, image_ref, (
                event.data if debug else ''))
        else:
            print('t', unixtime, event.ns, 'NO_IMAGE_REF', (
                event.data if (event.ns != 'moby' or debug) else ''))


def formatwarning(message, category, filename, lineno, line=None):
    """
    Override default Warning layout, from:

        /PATH/TO/ctre:123: UserWarning:
            Dupe ref found for image sha256:113...
          warn(

    To:

        ctre:123: UserWarning: Dupe ref found for image sha256:abc...
    """
    return '{basename}:{lineno}: {category}: {message}\n'.format(
        basename=basename(filename), lineno=lineno,
        category=category.__name__, message=message)
warnings.formatwarning = formatwarning  # noqa


def main():
    if len(argv) in (2, 3) and argv[1] == 'prune':
        force = (len(argv) == 3 and argv[2] in ('-f', '--force'))
        res = prunable_images_after_double_check()
        if res:
            if not force:
                print('NS', 'REF_TO_PRUNE')
                for ns, image in res:
                    print(ns, image)
                yn = ''
                while not yn.strip() or yn.strip() not in 'YyNn':
                    yn = input('Prune (y/n)? ')
                if yn.lower() not in 'Yy':
                    exit(1)

            # Remove the curated images.
            remove_images(res)

            # We could do the following even when !res, but we choose to only
            # do them when we appear to have more work anyway:
            # - Enumerate and try to remove all containers. Those that are
            #   running cannot be removed.
            remove_containers()
            # - Prune the content. Not entirely sure what this does, but it
            #   does clean up stuff that is unused.
            prune_content_references()

    elif len(argv) == 2 and argv[1] == 'record-image-starts':
        # Listen for events, check for container-add, lookup container image,
        # store last-used-time.
        record_image_starts()
    else:
        raise ValueError('please use: ctre prune [-f]')


if __name__ == '__main__':
    main()