tf-third-party/populate_cache.py at master · nvntsin/tf-third-party · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
#!/usr/bin/env python
#
# This script creates a local mirror for third-party packages used
# in Contrail builds.
#
# Usage: populate_cache.py target_directory/

import logging
import argparse
import os
import hashlib
import shutil
import sys
import urllib2

from urlparse import urlparse
from lxml import objectify

LOG = logging.getLogger("populate_cache")
LOG.setLevel(logging.DEBUG)

ch = logging.StreamHandler()
ch.setLevel(logging.DEBUG)

LOG.addHandler(ch)

def parse_packages_xml(path):
    return objectify.parse(path)

def get_package_list(xml):
    root = xml.getroot()
    return root.findall("package")

def get_package_details(element):
    md5sum = element['md5']
    canonical_url = element['urls'].findall("url[@canonical='true']")
    assert len(canonical_url) == 1

    try:
        filename = element['local-filename'].text
    except AttributeError:
        parsed_url = urlparse(canonical_url[0].text)
        filename = os.path.basename(parsed_url.path)

    return (canonical_url[0].text, filename, md5sum)

def cached_file_exists(dest, filename, md5sum):
    parent_dir = os.path.join(dest, filename[0])
    local_path = os.path.join(parent_dir, filename)

    if not os.path.exists(local_path):
        LOG.info("File %s missing, will download", filename)
        return False

    with open(local_path, 'rb') as fh:
        calculated_md5 = hashlib.md5(fh.read()).hexdigest()
        if calculated_md5 != md5sum:
            LOG.warn("File %s found, but checksum mismatch "
                     "(found: '%s' expected: '%s')",
                     filename, calculated_md5, md5sum)
            return False

    LOG.info("File %s found, checksum matches", filename)
    return True

def download_package(canonical_url, dest, filename, md5sum):
    parent_dir = os.path.join(dest, filename[0])
    local_path = os.path.join(parent_dir, filename)
    temp_path = local_path + '.tmp'

    if not os.path.exists(parent_dir):
        os.makedirs(parent_dir)

    chunk = 16 * 1024
    req = urllib2.urlopen(canonical_url)
    with open(temp_path, 'wb') as fh:
        shutil.copyfileobj(req, fh, chunk)

    with open(temp_path, 'rb') as fh:
        calculated_md5 = hashlib.md5(fh.read()).hexdigest()
        if calculated_md5 != md5sum:
            LOG.error("File %s checksum error"
                     "(found: '%s' expected: '%s')",
                     filename, calculated_md5, md5sum)
            return False

    LOG.info("File %s downloaded, checksum matches", filename)

    if os.path.exists(local_path):
        os.unlink(local_path)
    os.rename(temp_path, local_path)

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("destination", nargs=1)
    args = parser.parse_args(sys.argv[1:])

    dest = args.destination[0]

    xml = parse_packages_xml("packages.xml")
    packages = get_package_list(xml)
    for package in packages:
        canonical_url, filename, md5sum = get_package_details(package)
        if cached_file_exists(dest, filename, md5sum):
            continue
        download_package(canonical_url, dest, filename, md5sum)

if __name__ == "__main__":
    main()