#!/usr/bin/env python
#
# Maintenance script that is supposed to run as a cron-daily job.
# Makes sure the registry never uses more than a certain percentage of local disk space.
# Tries to keep the last versions of every image when recycling.
#

import sys
import os
import json
import subprocess
import re
import time
import syslog

REGISTRY = "{{ groups['registry']|first }}:{{ registry.port }}"
REGISTRY_URL = "https://%s" % REGISTRY
REGISTRY_NAME = "whisk_docker_registry"
REGISTRY_VERSION = "registry:{{ registry.version }}"
THRESHOLD = {{ registry.recycle.threshold }}
STORAGE = "{{ registry_storage_dir }}"
TAGS_TO_KEEP = {{ registry.recycle.keeptags }}

def getDiskUsage(dir):
    cmd=['df', '--output=pcent', dir]
    output = subprocess.check_output(cmd)
    pattern = "[0-9]+"
    result = re.search(pattern, output).group()
    return result

def getImageRepos():
    syslog.syslog("Getting image repositories from URL %s" % REGISTRY_URL)
    cmd=['curl', '-s', '-k', REGISTRY_URL + "/v2/_catalog"]
    output = subprocess.check_output(cmd)
    repos = json.loads(output)["repositories"]
    return repos

def getImageTags(repo):
    syslog.syslog("Getting tags for repo %s" % repo)
    cmd=['curl', '-s', '-k', REGISTRY_URL + "/v2/" + repo + "/tags/list"]
    output = subprocess.check_output(cmd)
    tags = json.loads(output)["tags"]
    tags_desc = sorted(tags, reverse=True)
    return tags_desc

def pullImage(repo, tag):
    syslog.syslog("Pulling image %s:%s" % (repo,tag))
    cmd=['docker', 'pull', REGISTRY + "/" + repo + ":" + tag]
    output = subprocess.check_output(cmd)
    return output

def pushImage(repo, tag):
    syslog.syslog("Pushing image %s:%s" % (repo,tag))
    cmd=['docker', 'push', REGISTRY + "/" + repo + ":" + tag]
    output = subprocess.check_output(cmd)
    return output

def deleteImage(repo, tag):
    syslog.syslog("Deleting image %s:%s" % (repo,tag))
    cmd=['docker', 'rmi', '-f', REGISTRY + "/" + repo + ":" + tag]
    output = subprocess.check_output(cmd)
    return output

def killRegistry():
    syslog.syslog("Killing registry %s" % REGISTRY_NAME)
    cmd=['docker', 'rm', '-f', '-v', REGISTRY_NAME]
    output = subprocess.check_output(cmd)
    return output

def startRegistry():
    syslog.syslog("Starting registry %s" % REGISTRY_NAME)
    cmd=['docker', 'run', '-d', '-p', '{{ registry.port }}:5000', '--name', REGISTRY_NAME, '--restart={{ docker.restart.policy }}', '-v', '{{ registry.confdir }}/certs:/certs', '-e',
    'REGISTRY_HTTP_TLS_CERTIFICATE=/certs/whisk_docker_registry.crt', '-e', 'REGISTRY_HTTP_TLS_KEY=/certs/whisk_docker_registry.key', REGISTRY_VERSION]
    output = subprocess.check_output(cmd)
    return output

############################################
# BEGIN MAIN PROGRAM
############################################

syslog.syslog("BEGIN REGISTRY CHECK")

# Step 1. Get disk space usage
USAGE=getDiskUsage(STORAGE)
syslog.syslog("Usage on %s is %s%%" % (STORAGE, USAGE))

if int(USAGE) < THRESHOLD:
  syslog.syslog("Usage below %s%%. Exiting." % THRESHOLD)
  syslog.syslog("REGISTRY CHECK END")
  exit(0)

syslog.syslog("Usage at least %s%%. Cleaning registry." % THRESHOLD)
REPO_TAG = {}

# Step 2. Get all image repos
REPOSITORIES=getImageRepos()

# Step 3. For each repo get tags
for repo in REPOSITORIES:
    REPO_TAG[repo] = []
    TAGS = getImageTags(repo)
    # Step 4. Pull last TAGS_TO_KEEP tags
    for i in range(TAGS_TO_KEEP):
        if len(TAGS) > i:
            tag = TAGS[i]
            pullImage(repo, tag)
            REPO_TAG[repo].append(tag)

# Step 5. Recycle registry
killRegistry()
startRegistry()

syslog.syslog("Waiting a couple of seconds for the registry to come up...")
time.sleep(5)

# Step 6. Push saved images back to new registry
for repo in REPOSITORIES:
    for tag in REPO_TAG[repo]:
        pushImage(repo, tag)

# Step 7. Delete pushed images locally
for repo in REPOSITORIES:
    for tag in REPO_TAG[repo]:
        deleteImage(repo, tag)

syslog.syslog("REGISTRY CHECK END")
