Merge "Fixed fetch_wheels.py to handle any docker-compliant registry."

scripts/fetch_wheels.py

@@ -2,25 +2,37 @@
 
 import json
 import os
+import re
+
 try:
     import urllib2
 except ImportError:
     # python3
     from urllib import request as urllib2
 
+DOCKER_REGISTRY='registry.hub.docker.com'
 
-def get_token(repo):
+def get_token(protocol, registry, repo):
-    url = "https://auth.docker.io/token?service=registry.docker.io&" \
+    if registry == DOCKER_REGISTRY:
-          "scope=repository:{}:pull".format(repo)
+      authserver = 'auth.docker.io'
+      service = 'registry.docker.io'
+    else:
+      authserver = "{}/v2".format(registry)
+      service =  registry.split(':')[0]
+    url = "{}://{}/token?service={}&" \
+            "scope=repository:{}:pull".format(protocol, authserver, service, repo)
+    print(url)
+    try:
+        r = urllib2.Request(url=url)
+        resp = urllib2.urlopen(r)
+        resp_text = resp.read().decode('utf-8').strip()
+        return json.loads(resp_text)['token']
+    except urllib2.HTTPError as err:
+        if err.reason == 'Not Found':
+            return None
 
-    r = urllib2.Request(url=url)
+def get_sha(repo, tag, registry, protocol, token):
-    resp = urllib2.urlopen(r)
+    url = "{}://{}/v2/{}/manifests/{}".format(protocol, registry, repo, tag)
-    resp_text = resp.read().decode('utf-8').strip()
-    return json.loads(resp_text)['token']
-
-
-def get_sha(repo, tag, registry, token):
-    url = "http://{}/v2/{}/manifests/{}".format(registry, repo, tag)
     print(url)
     r = urllib2.Request(url=url)
     if token:
@@ -30,9 +42,9 @@ def get_sha(repo, tag, registry, token):
     return json.loads(resp_text)['fsLayers'][0]['blobSum']
 
 
-def get_blob(repo, tag, registry='registry.hub.docker.com', token=None):
+def get_blob(repo, tag, protocol, registry=DOCKER_REGISTRY, token=None):
-    sha = get_sha(repo, tag, registry, token)
+    sha = get_sha(repo, tag, registry, protocol, token)
-    url = "http://{}/v2/{}/blobs/{} ".format(registry, repo, sha)
+    url = "{}://{}/v2/{}/blobs/{} ".format(protocol, registry, repo, sha)
     print(url)
     r = urllib2.Request(url=url)
     if token:
@@ -40,34 +52,45 @@ def get_blob(repo, tag, registry='registry.hub.docker.com', token=None):
     resp = urllib2.urlopen(r)
     return resp.read()
 
+def protocol_detection(registry, protocol='http'):
+    PROTOCOLS = ('http','https')
+    index = PROTOCOLS.index(protocol)
+    try:
+        url = "{}://{}".format(protocol, registry)
+        r = urllib2.Request(url)
+        resp = urllib2.urlopen(r)
+    except (urllib2.URLError,urllib2.HTTPError) as err:
+        if err.reason == 'Forbidden':
+            return protocol
+        elif index < len(PROTOCOLS) - 1:
+            return protocol_detection(registry, PROTOCOLS[index + 1])
+        else:
+            raise Exception, "Cannot detect protocol for registry: {} due to error: {}".format(registry,err)
+    except:
+        raise
+    else:
+        return protocol
 
 def get_wheels(url):
     r = urllib2.Request(url=url)
     resp = urllib2.urlopen(r)
     return resp.read()
 
-
 def parse_image(full_image):
-    if '/' in full_image:
+    slash_occurences = len(re.findall('/',full_image))
-        registry, image_with_tag = full_image.split('/', 1)
+    repo = None
+    registry = DOCKER_REGISTRY
+    if slash_occurences == 2:
+        registry, repo, image = full_image.split('/')
+    elif slash_occurences == 1:
+        repo, image = full_image.split('/')
     else:
-        registry = None
+        image = full_image
-        image_with_tag = full_image
+    if image.find(':') != -1:
-
+        image, tag = image.split(':')
-    if ':' in image_with_tag:
-        image, tag = image_with_tag.rsplit(':', 1)
     else:
-        image = image_with_tag
         tag = 'latest'
-
+    return registry, repo+'/'+image if repo else image, tag
-    if '/' in image:
-        return registry, image, tag
-
-    if registry:
-        return None, '/'.join([registry, image]), tag
-    else:
-        return None, image, tag
-
 
 def main():
     if 'WHEELS' in os.environ:
@@ -83,12 +106,12 @@ def main():
         data = get_wheels(wheels)
     else:
         registry, image, tag = parse_image(wheels)
+        protocol = protocol_detection(registry)
         kwargs = dict()
         if registry:
             kwargs.update({'registry': registry})
-        else:
+        kwargs.update({'token': get_token(protocol, registry, image)})
-            kwargs.update({'token': get_token(image)})
+        data = get_blob(image, tag, protocol, **kwargs)
-        data = get_blob(image, tag, **kwargs)
 
     if 'WHEELS_DEST' in os.environ:
         dest = os.environ['WHEELS_DEST']
@@ -102,3 +125,4 @@ def main():
 
 if __name__ == '__main__':
     main()
+