From 9b1daf3a5644318a25c4d1572c29b58023f01b64 Mon Sep 17 00:00:00 2001
From: Sam Yaple <sam@yaple.net>
Date: Sat, 30 Sep 2017 22:29:35 -0400
Subject: [PATCH] Add support to fetch_wheels for saved images

Change-Id: I622d75d3832c2c5312d6dbdb6c432100484dce17
---
 scripts/fetch_wheels.py | 51 +++++++++++++++++++++++++----------------
 1 file changed, 31 insertions(+), 20 deletions(-)

diff --git a/scripts/fetch_wheels.py b/scripts/fetch_wheels.py
index ceb0ef6..3be4191 100755
--- a/scripts/fetch_wheels.py
+++ b/scripts/fetch_wheels.py
@@ -68,26 +68,37 @@ def parse_image(full_image):
     else:
         return None, image, tag
 
-if 'WHEELS' in os.environ:
-    wheels = os.environ['WHEELS']
-else:
-    with open('/opt/loci/wheels', 'ro') as f:
-        wheels = f.read()
 
-with open('/opt/loci/wheels', 'w+') as f:
-    f.write(wheels)
-
-if wheels.startswith('http'):
-    data = get_wheels(wheels)
-else:
-    registry, image, tag = parse_image(wheels)
-    kwargs = dict()
-    if registry:
-        kwargs.update({'registry': registry})
+def main():
+    if 'WHEELS' in os.environ:
+        wheels = os.environ['WHEELS']
     else:
-        kwargs.update({'token': get_token(image)})
-    print(kwargs)
-    data = get_blob(image, tag, **kwargs)
+        with open('/opt/loci/wheels', 'ro') as f:
+            wheels = f.read()
 
-with open('/tmp/wheels.tar.gz', 'wb') as f:
-    f.write(data)
+    if wheels.startswith('/'):
+        with open(wheels, 'r') as f:
+            data = f.read()
+    elif wheels.startswith('http'):
+        data = get_wheels(wheels)
+    else:
+        registry, image, tag = parse_image(wheels)
+        kwargs = dict()
+        if registry:
+            kwargs.update({'registry': registry})
+        else:
+            kwargs.update({'token': get_token(image)})
+        data = get_blob(image, tag, **kwargs)
+
+    if 'WHEELS_DEST' in os.environ:
+        dest = os.environ['WHEELS_DEST']
+    else:
+        with open('/opt/loci/wheels', 'w') as f:
+            f.write(wheels)
+        dest = '/tmp/wheels.tar.gz'
+    with open(dest, 'wb') as f:
+        f.write(data)
+
+
+if __name__ == '__main__':
+    main()