Add argument parsing, ur_getenv, ur_check_version. Fix signal handling. Fix incorrect indexing of tarball progress output.

1 files changed, 114 insertions, 24 deletions

diff --git a/ur_upgrade.py b/ur_upgrade.py
index eada7ea..a22ae95 100755
--- a/ur_upgrade.py
+++ b/ur_upgrade.py
@@ -1,15 +1,19 @@
 #!/usr/bin/env python
 from __future__ import division
+
+import argparse
 import os
 import shutil
 import sys
-import stat
 import subprocess
 import tempfile
 import urllib2
 import signal
 import tarfile
+from distutils.version import StrictVersion
 from collections import namedtuple
+from string import Template
+from itertools import chain
 
 def which(cmd, mode=os.F_OK | os.X_OK, path=None):
     """Given a command, mode, and a PATH string, return the path which
@@ -74,6 +78,8 @@ def which(cmd, mode=os.F_OK | os.X_OK, path=None):
     return None
 
 class Ureka(object):
+    ''' 
+    '''
     def __init__(self, basepath):
         self.path = basepath
         self.path_data = os.path.abspath(os.path.join(self.path, 'misc'))
@@ -88,10 +94,13 @@ class Ureka(object):
         data = []
         for key in self._files:
             path = os.path.join(self.path_data, key)
-            item = open(path, 'r').readline().strip()
+            item = None
+            try:
+                item = open(path, 'r').readline().strip()
+            except:
+                item = ''
             data.append(item)
             
-            
         ureka_info = namedtuple('ureka_info', self._files)
         return ureka_info._make(data)
     
@@ -105,17 +114,81 @@ class Ureka(object):
         for item in self.info._asdict().iteritems():
             yield item
 
-UPGRADE_TEMP = tempfile.mkdtemp(prefix="upgrade")
-ORIGINAL = os.path.abspath('/home/jhunk/Ureka')
+def ur_getenv(ur_dir):
+    ''' Evaluate environment variables produced by ur-setup-real
+    '''
+    path = os.path.join(ur_dir, 'bin')
+    command = '{}/ur-setup-real -sh'.format(path).split()
+    output = subprocess.check_output(command, shell=True, stderr=open(os.path.devnull))
+    output = output.split(os.linesep)
+    output_env = []
+    
+    # Generate environment keypairs
+    for line in output:
+        if not line:
+            continue
+        if line.startswith('export'):
+            continue
+        line = line.strip()
+        line = line.replace(' ', '')
+        line = line.replace(';', '')
+        line = line.replace('"', '')
+        output_env.append(line.partition("=")[::2])
+    
+    output_env_temp = dict(output_env)
+    
+    # Perform shell expansion of Ureka's environment variables
+    for k, v in output_env_temp.items():
+        template = Template(v)
+        v = template.safe_substitute(output_env_temp)
+        output_env_temp[k] = v
+    
+    # Merge Ureka's environment with existing system environment
+    output_env_temp = dict(chain(os.environ.items(), output_env_temp.items()))
+    
+    # Assign expanded variables
+    output_env = output_env_temp
+    
+    return output_env
+
+def ur_check_version(urobj, vers):
+    if StrictVersion(urobj['version']) > vers:
+        return False
+    return True
+
+def cleanup(sig, stack):
+    print('')
+    print('Received signal {}!'.format(sig))
+    print('Cleaning up...')
+    shutil.rmtree(UPGRADE_TEMP)
+    exit(sig)
     
-def cleanup():
-    print('cleaning up')
 
 if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description='Ureka Upgrade Utility')
+    parser.add_argument('UR_DIR', action='store', type=str, help='Absolute path to Ureka installation')
+    parser.add_argument('--latest', action='store_true', default=False, help='Upgrade to the latest (stable) version')
+    parser.add_argument('--request', type=str, help='Upgrade to a specific version')
+    parser.add_argument('--mirror', type=str, help='Use a Ureka download mirror')
+    args = parser.parse_args()
     
+    if args.latest and args.request:
+        print("--latest and --request are mutually exclusive options.")
+        exit(1)
+    
+    if args.latest:
+        print("--latest is not implemented yet.")
+        exit(1)
         
+    if not args.request:
+        print("No version requested. Use --request (e.g. 1.4.1)")
+        exit(1)
+    
+    UPGRADE_TEMP = tempfile.mkdtemp(prefix="upgrade")
+    ORIGINAL = args.UR_DIR
+    
+    # Register signal handlers AFTER temporary directory is created
     signal.signal(signal.SIGINT, cleanup)
-    #signal.signal(signal.SIGKILL, cleanup)
     signal.signal(signal.SIGTERM, cleanup)
              
     if not which('rsync'):
@@ -123,15 +196,28 @@ if __name__ == "__main__":
         exit(1)
     
     u = Ureka(ORIGINAL)
-    INSTALL_PATH = os.path.join(UPGRADE_TEMP, '1.4.1')
-    UREKA_TARBALL = "Ureka_{}_{}_{}.tar.gz".format(u['os'], u['bits'], '1.4.1')
+    
+    if not ur_check_version(u, args.request):
+        print("Refusing to downgrade from {} to {}".format(u['version'], args.request))
+        exit(1)
+    
+    INSTALL_PATH = os.path.join(UPGRADE_TEMP, args.request)
+    UREKA_TARBALL = "Ureka_{}_{}_{}.tar.gz".format(u['os'], u['bits'], args.request)
     ''' http://ssb.stsci.edu/ureka/1.0/Ureka_linux-rhe5_64_1.4.1.tar.gz '''
 
-    #url = "http://ssb.stsci.edu/ureka/{}/Ureka_{}_{}_{}.tar.gz".format('1.4.1', u['os'], u['bits'], '1.4.1')
-    url = "http://localhost/ureka/{}/{}".format('1.4.1', UREKA_TARBALL)
+    #url = "http://ssb.stsci.edu/ureka/{}/{}".format(args.version, UREKA_TARBALL)
+    url = "http://localhost/ureka/{}/{}".format(args.request, UREKA_TARBALL)
+    
+    if args.mirror:
+        url = "{}/{}/{}".format(args.mirror, args.request, UREKA_TARBALL)
     
-    print("Downloading update... {}".format(url))
-    req_data = urllib2.urlopen(url)
+    print("Downloading archive... {}".format(url))
+    try:
+        req_data = urllib2.urlopen(url)
+    except Exception as ex:
+        print("Error {}, {}: {}".format(ex.code, ex.msg, url))
+        exit(1)
+        
     remote_size = float(req_data.headers['content-length'])
     total = 0
     with open(os.path.join(UPGRADE_TEMP, UREKA_TARBALL), 'wb') as installer:
@@ -143,10 +229,11 @@ if __name__ == "__main__":
     
     os.mkdir(INSTALL_PATH)
     tarball = tarfile.open(os.path.join(UPGRADE_TEMP, UREKA_TARBALL), 'r')
-    print("Preparing to unpack...")
-    files_total = 74690 #len(tarball.getmembers()) 
-    print("Unpacking update...")
-    for index, member in enumerate(tarball, start=0):
+    print("Preparing to unpack upgrade... (please wait)")
+    files_total = len(tarball.getmembers()) 
+    
+    print("Unpacking upgrade...")
+    for index, member in enumerate(tarball, start=1):
         print("\r{:.2f}% [{}/{}]".format((index / files_total * 100), index, files_total)),
         tarball.extract(member, path=INSTALL_PATH)
     print("")
@@ -154,18 +241,21 @@ if __name__ == "__main__":
     open(os.path.join(INSTALL_PATH, 'Ureka', 'misc', 'name'), 'w+').write(u['name'] + os.linesep)
     n = Ureka(os.path.join(INSTALL_PATH, 'Ureka'))
     
-    print("Updating distribution from {} to {}...".format(u['version'], n['version']))
+    print("Upgrading distribution {} to {}...".format(u['version'], n['version']))
     proc = subprocess.Popen('rsync -a -u {} {}'.format(n.path, os.path.dirname(u.path)).split())
     proc.wait()
     
-    # Regenerate Ureka object
+    # Regenerate upgraded ureka installation object
     u = Ureka(ORIGINAL)
 
-    print("Normalizing...")
-    proc = subprocess.Popen('{}/bin/ur-setup-real -sh'.format(u.path).split())
-    proc = subprocess.Popen('ur_normalize'.split())
+    print("Performing upgrade maintenance tasks...")
+    ur_env = ur_getenv(u.path)
     
-
+    proc = subprocess.Popen('ur_normalize -n -i -x'.split(), env=ur_env)
+    proc.wait()
     
+    print("Purging temporary data...")
     shutil.rmtree(UPGRADE_TEMP)
+    
+    print("Upgrade complete!")