path: root/images/package_images.py

#!/usr/bin/env python3
#
# Copyright 2018 Ettus Research, a National Instruments Company
#
# SPDX-License-Identifier: GPL-3.0-or-later
#
"""
Package image files into image archive packages

Provides functions for packaging image files into image packages. Generate the intermediate files
(like hash files), and create image archives from sets.
"""
from __future__ import print_function
import argparse
import copy
import glob
import hashlib
import itertools
import os
import re
import sys
import tempfile
import zipfile
from image_package_mapping import PACKAGE_MAPPING


def parse_args():
    """Setup argument parser and parse"""
    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument('--md5', action="store_true", default=False,
                        help="Generate MD5 files")
    parser.add_argument('--sha256', action="store_true", default=False,
                        help="Generate SHA256 files")
    parser.add_argument('-f', '--files', type=str, default="",
                        help="Comma separate list of files")
    parser.add_argument('-o', '--output', type=str, default="",
                        help="Output file to put the hashes in")
    parser.add_argument('-m', '--manifest', type=str, default="",
                        help="Update the manifest file at this path with the new SHAs")
    parser.add_argument('-t', '--targets', type=str, default="",
                        help="RegEx to select image sets from the manifest file.")
    parser.add_argument('-g', '--githash', type=str, default="",
                        help="Git hash directory name (eg. fpga-abc1234)")
    return parser.parse_args()


def gen_filelist(includes, excludes=None):
    """
    Generates a list of files, first generating
    :param includes: [list of] expression[s] to include
    :param excludes: [list of] expression[s] to exclude
    :return: flat list of filenames
    """
    if isinstance(includes, str):
        included = glob.glob(includes)
    else:
        included = list(itertools.chain(*[glob.iglob(filename) for filename in includes]))

    if excludes is None:
        excluded = []
    elif isinstance(excludes, str):
        excluded = glob.glob(excludes)
    else:
        excluded = list(itertools.chain(*[glob.iglob(filename) for filename in excludes]))
    # Remove the excluded files from the include list
    for filename in excluded:
        if filename in included:
            included.remove(filename)
    return included


def gen_md5(files_list, hash_filename=""):
    """Generate the .md5 files for all input files"""
    hashes = {}
    for filename in files_list:
        # Read and hash the input file
        with open(filename, 'rb') as img_file:
            md5_sum = hashlib.md5()
            md5_sum.update(img_file.read())
        # Write the hash to a *.md5 file
        with open(filename + '.md5', 'w') as md5_file:
            md5_hex = md5_sum.hexdigest()
            newline = "{md5_hex}  {filename}\n".format(filename=filename, md5_hex=md5_hex)
            md5_file.write(newline)
            # Also store it to write to a file of all the hashes
            hashes[filename] = md5_hex

    # Write the MD5 hashes to file
    with open(hash_filename, 'a') as hash_file:
        for filename, md5_hex in hashes.items():
            newline = "{md5_hex}  {filename}\n".format(filename=filename, md5_hex=md5_hex)
            hash_file.write(newline)


def gen_sha256(files_list, hash_filename=None, manifest_fn="", repo_and_hash=""):
    """Generate the SHA256 files for all input file"""
    # Input checking
    if hash_filename is None:
        hash_filename = "hashes.txt"
    print("Generating SHA256 sums for:\n{}".format(
        "\n".join("--{}".format(sha_fn) for sha_fn in files_list)))

    # Make a dictionary to store the new SHA256 sums
    sha256_dict = {}
    with open(hash_filename, 'a') as hash_file:
        for filename in files_list:
            with open(filename, 'rb') as img_file:
                sha256_sum = hashlib.sha256()
                sha256_sum.update(img_file.read())
                sha256_str = sha256_sum.hexdigest()
                newline = "{sha_hex}  {filename}\n".format(filename=filename,
                                                           sha_hex=sha256_str)
                hash_file.write(newline)
                # Add the sha256 to the dictionary
                basename = os.path.basename(filename)
                sha256_dict[basename] = sha256_str

    # If there's a manifest file to edit, put the new information in
    if os.path.isfile(manifest_fn):
        edit_manifest(manifest_fn, repo_and_hash, sha256_dict)


def gen_zip(zip_filename, files_list):
    """Generate the zip file for a set of images"""
    try:
        with zipfile.ZipFile(zip_filename, 'w', zipfile.ZIP_DEFLATED) as zip_file:
            for filename in files_list:
                zip_file.write(filename)
        return True
    except Exception as ex:
        print("Caught exception in gen_zip: {}".format(ex))
        return False


def do_gen_package(pkg_target, install_dir="", repo_and_hash=""):
    """Generate the entire N3XX image package, from the start to the end"""
    print("---Generating package for {}---".format(pkg_target))
    filelist = PACKAGE_MAPPING[pkg_target]['files']
    print("Required files:\n{}".format(
        "\n".join("--{}".format(img_fn) for img_fn in filelist)))
    md5_files = gen_filelist(includes=filelist, excludes=["*.rpt", "*.md5"])
    print("Files to md5sum:\n{}".format(
        "\n".join("--{}".format(md5_fn) for md5_fn in md5_files)))
    gen_md5(md5_files, "md5_hashes.txt")

    # Determine the current Git hash (w/o the repository)
    githash_l = re.findall(r"[\d\w]+-([\d\w]{7,8})", repo_and_hash)
    githash = githash_l[0] if githash_l else ""

    zip_files = gen_filelist(includes=filelist)
    zip_filename = os.path.join(install_dir, PACKAGE_MAPPING[pkg_target]['package_name'])\
        .format(githash)
    print("Files to zip:\n{}".format(
        "\n".join("--{}".format(zip_fn) for zip_fn in zip_files)))
    if not gen_zip(zip_filename, zip_files):
        zip_filename = ""
    return zip_filename


def gen_package(pkg_targets=(), repo_and_hash="", manifest_fn=""):
    """Generate the entire image package, and place it in the proper directory structure"""
    # Make the cache/ directory if necessary
    cache_path = os.path.join(os.getcwd(), "cache")
    if not os.path.isdir(cache_path):
        os.mkdir(cache_path)

    sha_filenames = []
    for pkg_target in pkg_targets:
        if pkg_target in PACKAGE_MAPPING:
            # Make the type directory
            pkg_type = PACKAGE_MAPPING[pkg_target]["type"]
            type_path = os.path.join(cache_path, pkg_type)
            if not os.path.isdir(type_path):
                os.mkdir(type_path)
            # Make the 'repository-hash' directory
            if not repo_and_hash:
                repo_and_hash = "repo-githash"
            git_path = os.path.join(type_path, repo_and_hash)
            if not os.path.isdir(git_path):
                os.mkdir(git_path)

            # Generate the package and add the the zip filename to the SHA list
            sha_filenames.append(do_gen_package(pkg_target,
                                                install_dir=git_path,
                                                repo_and_hash=repo_and_hash))
        else:
            print("Error: Specify a supported type from {}".format(
                list(PACKAGE_MAPPING.keys())))
    sha_filenames[:] = [sha_fn for sha_fn in sha_filenames if os.path.exists(sha_fn)]
    gen_sha256(sha_filenames, hash_filename="hashes.txt",
               manifest_fn=manifest_fn, repo_and_hash=repo_and_hash)


def list_differences(list1, list2):
    """Returns two lists containing the unique elements of each input list"""
    outlist1 = []
    outlist2 = []
    outlist1[:] = [elem for elem in list1 if elem not in list2]
    outlist2[:] = [elem for elem in list2 if elem not in list1]
    return outlist1, outlist2


def verify_package(zip_filename, pkg_target):
    """Verify the contents of the image package match the expected list of files"""
    expected_filelist = PACKAGE_MAPPING[pkg_target]['files']
    with zipfile.ZipFile(zip_filename, 'r') as zip_file:
        actual_filelist = zip_file.namelist()

    missing, extra = list_differences(expected_filelist, actual_filelist)
    if missing or extra:
        print("Error: image package does not include expected files")
        if missing:
            print("Missing files: {}".format(missing))
        if extra:
            print("Extra files: {}".format(extra))
        return False
    return True


def edit_manifest_line(line, new_repo_and_hash, new_hashes_dict):
    """Edit the line in the manifest to (maybe) include the new repo, git hash, and SHA"""
    # Check each value in your dictionary of new hashes
    for filename, new_hash in new_hashes_dict.items():
        # If the filename with a new hash shows up in the line
        if filename in line:
            # Replace the repo and git hash
            repo_and_hash = re.findall(r"[\w]+-[\da-fA-F]{7}", line)
            if repo_and_hash:
                repo_and_hash = repo_and_hash[0]
                line = line.replace(repo_and_hash, new_repo_and_hash)

            # Replace the SHA256
            sha = re.findall(r"[\da-fA-F]{64}", line)
            if sha:
                sha = sha[0]
                line = line.replace(sha, new_hash)

            if not repo_and_hash or not sha:
                print("Error: repo, hash or SHA missing in line with new file")
                print("Line: {}", format(line))
            # If we found and replaced info, return the edited line
            return line
    # If we never edit the line, just return it
    return line


def edit_manifest(manifest_fn, new_repo_and_hash, new_hash_dict):
    """Edit the provided manifest file to update the githash and SHA256"""
    with tempfile.NamedTemporaryFile(mode='w', dir='.', delete=False) as tmp_manifest, \
            open(manifest_fn, 'r') as old_manifest:
        # Check each line in the manifest file
        for line in old_manifest:
            # If needed, put the new info in the line
            line = edit_manifest_line(line, new_repo_and_hash, new_hash_dict)
            # Always write the line back
            tmp_manifest.write(line)
    # Replace the manifest file with our temporary file that we created
    os.rename(tmp_manifest.name, manifest_fn)


def determine_targets():
    """
    Determine which image packages can be created by the files in the current directory
    :return: list of valid targets
    """
    found_targets = []
    for target, target_info in PACKAGE_MAPPING.items():
        # Grab the list of files required, but remove any files that we're going to build here,
        # like the hash files
        required_files = copy.deepcopy(target_info['files'])
        required_files[:] = [filename for filename in required_files if '.md5' not in filename]

        if all([os.path.exists(img_file) for img_file in required_files]):
            found_targets.append(target)
    return found_targets


def main():
    """Generate image packages using commandline arguments"""
    args = parse_args()
    if args.md5 or args.sha256 or args.files or args.output:
        print("Unsupported argument: only --pkg_targets is currently supported.")
    if args.targets:
        pkg_targets = [ss.strip() for ss in args.targets.split(',')]
    else:
        pkg_targets = determine_targets()
        print("Targets to package:\n{}".format(
            "\n".join("--{}".format(pkg) for pkg in pkg_targets)))
    gen_package(pkg_targets=pkg_targets, repo_and_hash=args.githash, manifest_fn=args.manifest)
    return True


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