move json loading and stashserver loading into functions

Author: DrDaveD

stashcp/__init__.py

@@ -563,6 +563,138 @@ def get_ips(name):
     # always prefer IPv4
     return ipv4s + ipv6s
 
+
+# Return list of cache URLs
+def get_json_caches(caches_json_location):
+    try:
+        with open(caches_json_location, 'r') as f:
+            caches_list = json.loads(f.read())
+            logging.debug("Loaded caches list from %s", caches_json_location)
+    except:
+        logging.error("Unable to open or parse json in %s: %s", 
+            caches_json_location, str(sys.exc_info()[1]))
+        return None
+
+    usable_caches = []
+    for cache in caches_list:
+        if 'status' in cache and cache['status'] == 0:
+            continue
+        if 'name' in cache:
+            usable_caches.append(cache['name'])
+    if len(usable_caches) == 0:
+        logging.error("No cache names found in %s without zero status", caches_json_location)
+        return None
+
+    return usable_caches
+
+
+# Return list of caches as root:// URLs
+def get_stashservers_caches(responselines):
+
+    # After the geo order of the selected server list on line zero,
+    #  the rest of the response is in .cvmfswhitelist format.
+    # This is done to avoid using https for every request on the
+    #  wlcg-wpad servers and takes advantage of conveniently
+    #  existing infrastructure.
+    # The format contains the following lines:
+    # 1. Creation date stamp, e.g. 20200414170005.  For debugging
+    #    only.
+    # 2. Expiration date stamp, e.g. E20200421170005.  cvmfs clients
+    #    check this to avoid replay attacks, but for this api that
+    #    is not much of a risk so it is ignored.
+    # 3. "Repository" name, e.g. Nstash-servers.  cvmfs clients
+    #    also check this but it is not important here.
+    # 4. With cvmfs the 4th line has a repository fingerprint, but
+    #    for this api it instead contains a semi-colon separated list
+    #    of named server lists.  Each server list is of the form
+    #    name=servers where servers is comma-separated.  Ends with
+    #    "hash=-sha1" because cvmfs_server expects the hash name
+    #    to be there.  e.g.
+    #    xroot=stashcache.t2.ucsd.edu,sg-gftp.pace.gatech.edu;xroots=xrootd-local.unl.edu,stashcache.t2.ucsd.edu;hash=-sha1
+    # 5. A two-dash separator, i.e "--"
+    # 6. The sha1 hash of lines 1 through 4.
+    # 7. The signature, i.e. an RSA encryption of the hash that can
+    #    be decrypted by the OSG cvmfs public key.  Contains binary
+    #    information so it may contain a variable number of newlines
+    #    which would have caused it to have been split into multiple
+    #    response "lines".
+
+    if len(responselines) < 8:
+        logging.error("stashservers response too short, less than 8 lines")
+        return None
+    hashname = responselines[4][-5:]
+    if hashname != "-sha1":
+        logging.error("stashservers response does not have sha1 hash: %s", hashname)
+        return None
+    hashedtext = '\n'.join(responselines[1:5]) + '\n'
+    hash = hashlib.sha1(hashedtext).hexdigest()
+    if responselines[6] != hash:
+        logging.debug("stashservers hash %s does not match expected hash %s", responselines[6], hash)
+        logging.debug("hashed text:\n%s", hashedtext)
+        logging.error("stashservers response hash does not match expected hash")
+        return None
+
+    # Call out to /usr/bin/openssl if present, in order to avoid
+    #  python dependency on a crypto package.
+    if not os.path.exists("/usr/bin/openssl"):
+        # The signature check isn't critical to be done everywhere;
+        #  any tampering will likely to be caught somewhere and
+        #  investigated.  Usually openssl is present.
+        logging.debug("openssl not installed, skipping signature check")
+    else:
+        sig = '\n'.join(responselines[7:])
+
+        # Look for the OSG cvmfs public key to verify signature
+        prefix = os.environ.get("OSG_LOCATION", "/")
+        osgpub = 'opensciencegrid.org.pub'
+        pubkey_files = ['/etc/cvmfs/keys/opensciencegrid.org/' + osgpub,
+                        os.path.join(prefix, "etc/stashcache", osgpub),
+                        os.path.join(prefix, "usr/share/stashcache", osgpub)]
+        if resource_filename:
+            try:
+                pubkey_files.append(resource_filename(__name__, osgpub))
+            except IOError as ioe:
+                logging.debug("Unable to retrieve caches.json using resource string, trying other locations")
+
+        for pubkey_file in pubkey_files:
+            if os.path.isfile(pubkey_file):
+                break
+        else:
+            logging.error("Unable to find osg cvmfs key in %r", pubkey_files)
+            return None
+        
+        cmd = "/usr/bin/openssl rsautl -verify -pubin -inkey " + pubkey_file
+        logging.debug("Running %s", cmd)
+        p = subprocess.Popen(cmd, shell=True,
+                stdin=subprocess.PIPE, stdout=subprocess.PIPE)
+        p.stdin.write(sig)
+        p.stdin.close()
+        decryptedhash = p.stdout.read()
+        p.stdout.close()
+        if hash != decryptedhash:
+            logging.debug("stashservers hash %s does not match decrypted signature %s", hash, decryptedhash)
+            logging.error("stashservers signature does not verify")
+            return None
+        logging.debug("Signature matched")
+
+    lists = responselines[4].split(';')
+    logging.debug("Cache lists: %s", lists)
+
+    if cache_list_name == None:
+        caches = lists[0].split('=')[1]
+    else:
+        for l in lists:
+            n=len(cache_list_name)+1
+            if l[0:n] == cache_list_name + '=':
+                caches = l[n:]
+                break
+    caches_list = caches.split(',')
+    for i in range(len(caches_list)):
+        caches_list[i] = 'root://' + caches_list[i]
+    
+    return caches_list
+
+
 # Return best stashcache and set nearest_cache_list global
 def get_best_stashcache():
     global nearest_cache_list
@@ -576,48 +708,29 @@ def get_best_stashcache():
     # Randomize the geo ip sites
     random.shuffle(geo_ip_sites)
 
-    api_text = ''
+    api_text = ""
 
     caches_list = []
 
     # Check if the user provided a caches json file location
-    if caches_json_location and os.path.exists(caches_json_location):
-        # Use geo ip api on caches in provided json file
-        try:
-            with open(caches_json_location, 'r') as f:
-                caches_list = json.loads(f.read())
-                logging.debug("Loaded caches list from %s", caches_json_location)
-        except:
-            logging.error("Unable to open or parse json in %s: %s", 
-                caches_json_location, str(sys.exc_info()[1]))
+    if caches_json_location:
+        if not os.path.exists(caches_json_location):
+            logging.error(caches_json_location + " does not exist")
             return None
-
-        # Format the caches for the GeoIP query
+        # Use geo ip api on caches in provided json file
+        caches_list = get_json_caches(caches_json_location)
         caches_string = ""
-        usable_caches = []
         for cache in caches_list:
-            if 'status' in cache and cache['status'] == 0:
-                continue
-            if 'name' in cache:
-                usable_caches.append(cache['name'])
-                parsed_url = urlparse(cache['name'])
-                caches_string = "%s,%s" % (caches_string, parsed_url.hostname)
-        if len(usable_caches) == 0:
-            logging.error("No cache names found in %s without zero status", caches_json_location)
-            return None
-
-        caches_list = usable_caches
-
+            parsed_url = urlparse(cache)
+            caches_string = "%s,%s" % (caches_string, parsed_url.hostname)
         # Remove the first comma
         caches_string = caches_string[1:]
-        
         api_text = "api/v1.0/geo/stashcp/" + caches_string
-        
     else:
         # Use stashservers.dat api
         api_text = "stashservers.dat"
         if cache_list_name != None:
-            api_text += '?list=' + cache_list_name
+            api_text += "?list=" + cache_list_name
 
     responselines = []
     i = 0
@@ -644,11 +757,11 @@ def get_best_stashcache():
                 logging.debug("Error: %s", str(e))
         i+=1
 
-    order_str = ''
+    order_str = ""
     if len(responselines) > 0:
         order_str = responselines[0]
 
-    if order_str == '':
+    if order_str == "":
         if len(caches_list) == 0:
             logging.error("unable to get list of caches")
             return None
@@ -667,107 +780,9 @@ def get_best_stashcache():
 
         if len(caches_list) == 0:
             # Used the stashservers.dat api
-
-            # After the geo order of the selected server list on line zero,
-            #  the rest of the response is in .cvmfswhitelist format.
-            # This is done to avoid using https for every request on the
-            #  wlcg-wpad servers and takes advantage of conveniently
-            #  existing infrastructure.
-            # The format contains the following lines:
-            # 1. Creation date stamp, e.g. 20200414170005.  For debugging
-            #    only.
-            # 2. Expiration date stamp, e.g. E20200421170005.  cvmfs clients
-            #    check this to avoid replay attacks, but for this api that
-            #    is not much of a risk so it is ignored.
-            # 3. "Repository" name, e.g. Nstash-servers.  cvmfs clients
-            #    also check this but it is not important here.
-            # 4. With cvmfs the 4th line has a repository fingerprint, but
-            #    for this api it instead contains a semi-colon separated list
-            #    of named server lists.  Each server list is of the form
-            #    name=servers where servers is comma-separated.  Ends with
-            #    "hash=-sha1" because cvmfs_server expects the hash name
-            #    to be there.  e.g.
-            #    xroot=stashcache.t2.ucsd.edu,sg-gftp.pace.gatech.edu;xroots=xrootd-local.unl.edu,stashcache.t2.ucsd.edu;hash=-sha1
-            # 5. A two-dash separator, i.e "--"
-            # 6. The sha1 hash of lines 1 through 4.
-            # 7. The signature, i.e. an RSA encryption of the hash that can
-            #    be decrypted by the OSG cvmfs public key.  Contains binary
-            #    information so it may contain a variable number of newlines
-            #    which would have caused it to have been split into multiple
-            #    response "lines".
-
-            if len(responselines) < 8:
-                logging.error("stashservers response too short, less than 8 lines")
-                return None
-            hashname = responselines[4][-5:]
-            if hashname != "-sha1":
-                logging.error("stashservers response does not have sha1 hash: %s", hashname)
+            caches_list = get_stashservers_caches(responselines)
+            if caches_list is None:
                 return None
-            hashedtext = '\n'.join(responselines[1:5]) + '\n'
-            hash = hashlib.sha1(hashedtext).hexdigest()
-            if responselines[6] != hash:
-                logging.debug("stashservers hash %s does not match expected hash %s", responselines[6], hash)
-                logging.debug("hashed text:\n%s", hashedtext)
-                logging.error("stashservers response hash does not match expected hash")
-                return None
-
-            # Call out to /usr/bin/openssl if present, in order to avoid
-            #  python dependency on a crypto package.
-            if not os.path.exists("/usr/bin/openssl"):
-                # The signature check isn't critical to be done everywhere;
-                #  any tampering will likely to be caught somewhere and
-                #  investigated.  Usually openssl is present.
-                logging.debug("openssl not installed, skipping signature check")
-            else:
-                sig = '\n'.join(responselines[7:])
-
-                # Look for the OSG cvmfs public key to verify signature
-                prefix = os.environ.get("OSG_LOCATION", "/")
-                osgpub = 'opensciencegrid.org.pub'
-                pubkey_files = ['/etc/cvmfs/keys/opensciencegrid.org/' + osgpub,
-                                os.path.join(prefix, "etc/stashcache", osgpub),
-                                os.path.join(prefix, "usr/share/stashcache", osgpub)]
-                if resource_filename:
-                    try:
-                        pubkey_files.append(resource_filename(__name__, osgpub))
-                    except IOError as ioe:
-                        logging.debug("Unable to retrieve caches.json using resource string, trying other locations")
-
-                for pubkey_file in pubkey_files:
-                    if os.path.isfile(pubkey_file):
-                        break
-                else:
-                    logging.error("Unable to find osg cvmfs key in %r", pubkey_files)
-                    return None
-                
-                cmd = "/usr/bin/openssl rsautl -verify -pubin -inkey " + pubkey_file
-                logging.debug("Running %s", cmd)
-                p = subprocess.Popen(cmd, shell=True,
-                        stdin=subprocess.PIPE, stdout=subprocess.PIPE)
-                p.stdin.write(sig)
-                p.stdin.close()
-                decryptedhash = p.stdout.read()
-                p.stdout.close()
-                if hash != decryptedhash:
-                    logging.debug("stashservers hash %s does not match decrypted signature %s", hash, decryptedhash)
-                    logging.error("stashservers signature does not verify")
-                    return None
-                logging.debug("Signature matched")
-
-            lists = responselines[4].split(';')
-            logging.debug("Cache lists: %s", lists)
-
-            if cache_list_name == None:
-                caches = lists[0].split('=')[1]
-            else:
-                for l in lists:
-                    n=len(cache_list_name)+1
-                    if l[0:n] == cache_list_name + '=':
-                        caches = l[n:]
-                        break
-            caches_list = caches.split(',')
-            for i in range(len(caches_list)):
-                caches_list[i] = 'root://' + caches_list[i]
 
         minsite = caches_list[int(ordered_list[0])-1]