From ed1b190e478e6e050db240ffb1f2a8107c86786d Mon Sep 17 00:00:00 2001
From: Aloys Baillet <aloys.baillet+github@gmail.com>
Date: Wed, 19 Aug 2020 17:23:39 +1000
Subject: [PATCH] Added dockerstats command

Signed-off-by: Aloys Baillet <aloys.baillet+github@gmail.com>
---
 python/aswfdocker/cli/aswfdocker.py | 25 +++++++++++++++++
 python/aswfdocker/utils.py          | 42 +++++++++++++++++++++++++++++
 2 files changed, 67 insertions(+)

diff --git a/python/aswfdocker/cli/aswfdocker.py b/python/aswfdocker/cli/aswfdocker.py
index fe184d00..c77a8bac 100644
--- a/python/aswfdocker/cli/aswfdocker.py
+++ b/python/aswfdocker/cli/aswfdocker.py
@@ -267,6 +267,31 @@ def images():
                     click.echo(f"{group}/{image_name}:{version}")
 
 
+@click.option(
+    "--sizes", "-s", is_flag=True, help="Display Sizes",
+)
+@cli.command()
+def dockerstats(sizes):
+    """Lists all known ci images in this format: IMAGEGROUP/IMAGE:VERSION
+    """
+    if sizes:
+        total_size = 0
+        for org, image_type, image in utils.iter_all_images():
+            image_name = utils.get_image_name(image_type, image)
+            for tag, size in utils.get_image_sizes(org, image_name).items():
+                click.echo(f"{org}/{image_name}:{tag} size={size}")
+                total_size += size
+        click.echo(f"Total size={total_size}")
+    else:
+        total_pull = 0
+        for org, image_type, image in utils.iter_all_images():
+            image_name = utils.get_image_name(image_type, image)
+            pull_count = utils.get_image_pull_count(org, image_name)
+            click.echo(f"{org}/{image_name} pull_count={pull_count}")
+            total_pull += pull_count
+        click.echo(f"Total pull_count={total_pull}")
+
+
 @cli.command()
 @click.option(
     "--settings-path", "-p", default="~/.aswfdocker", help="User settings file path.",
diff --git a/python/aswfdocker/utils.py b/python/aswfdocker/utils.py
index 029e8077..e41af520 100644
--- a/python/aswfdocker/utils.py
+++ b/python/aswfdocker/utils.py
@@ -8,6 +8,8 @@
 import subprocess
 import datetime
 import logging
+import json
+import urllib.request
 
 from aswfdocker import constants
 
@@ -119,3 +121,43 @@ def get_group_from_image(image_type: constants.ImageType, image: str):
             if img == image:
                 return group
     raise RuntimeError(f"Cannot find group for image {image}")
+
+
+def get_image_pull_count(docker_org, image):
+    url = f"https://hub.docker.com/v2/repositories/{docker_org}/{image}"
+    try:
+        d = json.loads(urllib.request.urlopen(url).read())
+        return d["pull_count"]
+    except urllib.error.HTTPError:
+        logger.debug("Failed to load data from URL %r", url)
+        return 0
+
+
+def get_image_sizes(docker_org, image):
+    sizes = {}
+    url = f"https://hub.docker.com/v2/repositories/{docker_org}/{image}/tags/"
+    try:
+        d = json.loads(urllib.request.urlopen(url).read())
+    except urllib.error.HTTPError:
+        logger.debug("Failed to load data from URL %r", url)
+        return sizes
+    digests = set()
+    for tag in d["results"]:
+        digest = tag["images"][0]["digest"]
+        if digest in digests:
+            continue
+        digests.add(digest)
+        sizes[tag["name"]] = tag["full_size"]
+    return sizes
+
+
+def iter_all_images():
+    for org in (constants.TESTING_DOCKER_ORG, constants.PUBLISH_DOCKER_ORG):
+        for image_type in (
+            constants.ImageType.PACKAGE,
+            constants.ImageType.CI_IMAGE,
+            constants.ImageType.RT_IMAGE,
+        ):
+            for _, images in constants.GROUPS[image_type].items():
+                for image in images:
+                    yield org, image_type, image