Quick and dirty test of functions to freeze a Dockerfile

test.py

#! /usr/bin/env python3

# Copyright (c) 2019 VMware, Inc. All Rights Reserved.
# SPDX-License-Identifier: BSD-2-Clause

from datetime import datetime
import os
from pprint import pprint
import sys

import docker
from dockerfile_parse import DockerfileParser


def is_universal_os():
    return False

# TODO: test for standard base image, multistage with same base, multistage with
# differing base all with and without ARG to define base
def get_base_image(dfp):
    base = dfp.baseimage
    from_arg = None
    # if dfp.baseimage is the value of an ARG, find out what the default is
    if base.startswith('$'):
        from_arg = dfp.baseimage.lstrip('$')
        for node in dfp.structure:
            if node['instruction'] == 'ARG' and node.get('value', '').startswith(from_arg):
                base = node['value'].split('=')[1]
                break
    return base, from_arg

def get_image_digest(image, tag, df):
    dock = docker.from_env()

    img = None
    try:
        # check we have an up-to-date image
        img = dock.images.pull(image, tag=tag)
    except errors.DockerException as err:
        print(err.message)
    if img:
        return img.id
        
def now_to_snapshot():
    now = datetime.utcnow()
    return now.isoformat(timespec='seconds').replace('-','').replace(':','')+'Z'

def set_apt_source_as_snapshot():
    return None

def handle_docker_file(df):
    if not os.path.exists(df):
        print("Can't open non-existent file")
        return

    print("Working with file %s" % df)
    dfp = DockerfileParser(fileobj=open(df, 'r'))
    pprint(dfp.structure)
    # TODO: how to serialise the new Dockerfile without overwriting the existing
    # one? Creating a DockerfileParser() without args seems to result in writes
    # to the fileobj in dfp, even though it's opened with only 'r' permissions
    # there
    newdf = []

    # TODO: if base image was from an ARG we need to delete that ARG in newdf
    base, from_arg = get_base_image(dfp)
    digest = None
    image = base
    if image.__contains__('@'):
        image, digest = base.split('@', 1)
    image, tag = base.split(':', 1)
    if not digest:
        digest = get_image_digest(image, tag, df).split(':', 1)[1]
    # print("Base image of container is '%s'" % image)
    # print("Using image tag of '%s'" % tag)
    # print("Using digest of '%s'" % digest)

    newdf.append('FROM %s:%s@%s\n' % (image, tag, digest))
    for node in dfp.structure:
        print("Node of type %s with value: '%s'" % (node['instruction'], node['value']))
        if node['instruction'] == 'FROM':
            # TODO: could have multiple FROM statements
            # print(node['content'])
            continue
        if node['instruction'] == 'ARG' and node.get('value', '').startswith(from_arg):
            continue
        newdf.append(node['content'])
    # print(newdf)
    print('\n'.join(newdf))


def run_dff_test():
    if len(sys.argv) < 2:
        print("Need at least one arg, a Dockerfile")
        return

    handle_docker_file(sys.argv[1])


if __name__ == "__main__":
    run_dff_test()