source: trunk/bin/asap_update_data

#!/usr/bin/env python
from __future__ import print_function
import os
import sys
import tempfile
import urllib
import errno
import shutil
import tarfile

def errormsg(args, **kwargs):
    print(*args, file=sys.stderr, **kwargs)

def isWritable(path):
    try:
        testfile = tempfile.TemporaryFile(dir = path)
        testfile.close()
    except OSError as e:
        if e.errno == errno.EACCES:  # 13
            return False
        e.filename = path
        raise
    return True

def cleanup(inlist):
    for item in inlist:
        if os.path.isfile(item):
            os.remove(item)
        if os.path.isdir(item):
            os.rmdir(item)

# globals
md5suff  = '.md5sum'
name     = 'asap_data.tar.bz2'
dataurl  = "ftp://ftp.atnf.csiro.au"
datadir  = 'pub/software/asap/data'

# use ASAPDATA if set - allows non-root update
if os.environ.has_key("ASAPDATA"):
    asapbase = os.environ["ASAPDATA"]
else:
    import asap
    # get asap module path
    asapbase = asap.__path__[0]

print('asapbase is ' + asapbase)
# check we have write permission to asapbase before continuing
if not isWritable(asapbase):
    errormsg(['No write access to %s, aborting ' % asapbase])
    exit()

# Ensure the file is read/write by the creator only
saved_umask = os.umask(0077)

tmpdir  = tempfile.mkdtemp()
md5path = os.path.join(tmpdir, name + md5suff)

print("Checking if an update is required.")
url = dataurl + '/' + datadir + '/' + name + md5suff
try:
    urllib.urlretrieve (url, md5path)
    md5file = file(md5path)
    md5new = md5file.readlines()[0].split()[0]
except IOError as e:
    errormsg(['Download failed\n','IOError %s' % e.errno])
    cleanup([md5path,tmpdir])
    exit()
else:
    print('Downloaded checksum file to ' + md5path)

try:
    fl = os.path.join(asapbase, name+md5suff)
    data_md5 = file(fl)
    ls = data_md5.readlines()[0]
    data_md5.close()
    md5old = ls.split()[0]
except IOError:
    md5old =''

if md5new == md5old:
    print("""Data already at latest available version.
If you still get errors running asap, please report this.""")
    cleanup([md5path,tmpdir])
    exit()

print("Update required. Downloading asap data archive....")
url = dataurl + '/' + datadir + '/' + name
tarpath = os.path.join(tmpdir, name)
urllib.urlcleanup()
try:
     urllib.urlretrieve (url, tarpath)
except IOError as e:
    errormsg(['Download failed\n','IOError %s' % e.errno])
    cleanup([tarpath,md5path,tmpdir])
    exit()
else:
    print('Downloaded tar file to ' + tarpath)

print("Extracting data archive in %s.." % asapbase)
os.umask(saved_umask)
shutil.copy(md5path, fl)
shutil.copy(tarpath, os.path.join(asapbase, name))
os.chdir(asapbase)
tf = tarfile.TarFile.bz2open(name)
for member in tf.getmembers():
    tf.extract(member)

cleanup([tarpath,md5path,tmpdir])
exit()