xnat_sync

#!/usr/bin/python

# Copyright (c) 2011 Christian Haselgrove
# BSD License: http://www.opensource.org/licenses/bsd-license.php

import sys
import os
import time
import datetime
import getpass
import dateutil.parser
import pyxnat

def command_line_error(msg):
    sys.stderr.write('%s: %s\n' % (progname, msg))
    sys.stderr.write('run %s with no arguments for usage\n' % progname)
    return

def write_rc(xnat_uri, resource_uri, files):
    fo = open('.xnat', 'w')
    fo.write('%s\n' % xnat_uri)
    fo.write('%s\n' % resource_uri)
    fo.write('%s\n' % datetime.datetime.now().strftime('%s'))
    for path in sorted(files):
        times = files[path]
        fo.write('%d %d %s\n' % (times['t_local'], times['t_remote'], path))
    fo.close()
    return

def read_rc():
    fo = open('.xnat')
    xnat_uri = fo.readline().strip()
    resource_uri = fo.readline().strip()
    last_updated = int(fo.readline().strip())
    files = {}
    for line in fo:
        (t_local, t_remote, path) = line.strip('\n').split(' ', 2)
        files[path] = {'t_remote': int(t_remote), 't_local': int(t_local)}
    fo.close()
    return (xnat_uri, resource_uri, last_updated, files)

def get_statuses(rc_files, resource):
    all_files = set(rc_files)
    remote_files = {}
    for f in resource.files():
        path = f.attributes()['path']
        t_remote = int(dateutil.parser.parse(f.last_modified()).strftime('%s'))
        remote_files[path] = t_remote
        all_files.add(path)
    local_files = {}
    for (dirpath, dirnames, filenames) in os.walk('.'):
        for filename in filenames:
            path = '%s/%s' % (dirpath, filename)
            # remote initial './'
            path = path[2:]
            if path == '.xnat':
                continue
            local_files[path] = int(os.stat(path).st_mtime)
            all_files.add(path)
    statuses = {}
    for path in sorted(all_files):
        if path not in rc_files:
            if path in local_files:
                local_code = 'N'
            else:
                local_code = '-'
        elif path not in local_files:
            local_code = 'D'
        elif local_files[path] == rc_files[path]['t_local']:
            local_code = 'U'
        elif local_files[path] > rc_files[path]['t_local']:
            local_code = 'M'
        else:
            local_code = '?'
        if path not in rc_files:
            if path in remote_files:
                remote_code = 'N'
            else:
                remote_code = '-'
        elif path not in remote_files:
            remote_code = 'D'
        elif remote_files[path] == rc_files[path]['t_remote']:
            remote_code = 'U'
        elif remote_files[path] > rc_files[path]['t_remote']:
            remote_code = 'M'
        else:
            remote_code = '?'
        statuses[path] = {'local': local_code, 'remote': remote_code}
    return statuses

def files_times(path, f):
    t_local = int(os.stat(path).st_mtime)
    t_remote = int(dateutil.parser.parse(f.last_modified()).strftime('%s'))
    return {'t_remote': t_remote, 't_local': t_local}

def connect(uri, user, password):
    if user is None:
        return pyxnat.Interface(uri, anonymous=True)
    return pyxnat.Interface(uri, user, password)

def display_project_resources(project):
    print '+ Project %s' % project.id()
    display_resources(project.resources(), '    ')
    for subject in project.subjects():
        display_subject_resources(subject, '    ')
    return

def display_subject_resources(subject, prefix=''):
    print '%s+ Subject %s' % (prefix, subject.label())
    display_resources(subject.resources(), prefix)
    for experiment in subject.experiments():
        display_experiment_resources(experiment, prefix + '    ')
    return

def display_experiment_resources(experiment, prefix=''):
    print '%s+ Experiment %s' % (prefix, experiment.label())
    display_resources(experiment.resources(), prefix)
    for scan in experiment.scans():
        display_scan_resources(scan, prefix + '    ')
    for reconstruction in experiment.reconstructions():
        print '%s    + Reconstruction %s' % (prefix, reconstruction.label())
        display_resources(reconstruction.in_resources(), prefix + '    ', 'In')
        display_resources(reconstruction.out_resources(), prefix + '    ', 'Out')
    for assessor in experiment.assessors():
        print '%s    + Assessment %s' % (prefix, assessor.label())
        display_resources(assessor.in_resources(), prefix + '    ', 'In')
        display_resources(assessor.out_resources(), prefix + '    ', 'Out')
    return

def display_scan_resources(scan, prefix=''):
    print '%s+ Scan %s' % (prefix, scan.id())
    display_resources(scan.resources(), prefix + '    ')
    return

def display_resources(cobject, prefix, res_pre=None):
    resources = [ r for r in cobject ]
    resources.sort(lambda a, b: cmp(int(a.id()), int(b.id())))
    for r in resources:
        n_files = len([ f for f in r.files() ])
        if not res_pre:
            print '%s  * Resource %5s: %s (%d files)' % (prefix, r.id(), r.label(), n_files)
        else:
            print '%s  * %s Resource %5s: %s (%d files)' % (prefix, res_pre, r.id(), r.label(), n_files)

def get_object(i, argv):
    """return a pyxnat object given given command line:
        project [subject [experiment]]
    prints a message and returns None on error
    """
    if not sys.argv:
        command_line_error('no project given')
        return None
    project_id = sys.argv.pop(0)
    project = i.select.project(project_id)
    if not project.exists():
        sys.stderr.write("%s: can't find project %s\n" % (progname, project_id))
        return None
    if not sys.argv:
        return project
    subject_id = sys.argv.pop(0)
    subject = project.subject(subject_id)
    if not subject.exists():
        sys.stderr.write("%s: can't find subject %s for project %s\n" % (progname, subject_id, project_id))
        return None
    if not sys.argv:
        return subject
    experiment_id = sys.argv.pop(0)
    experiment = subject.experiment(experiment_id)
    if not experiment.exists():
        sys.stderr.write("%s: can't find experiment %s for subject %s\n" % (progname, experiment_id, subject_id))
        return None
    if not sys.argv:
        return experiment
    scan_name = sys.argv.pop(0)
    scan = experiment.scan(scan_name)
    if not scan.exists():
        sys.stderr.write("%s: can't find scan %s for experiment %s\n" % (progname, scan_name, experiment_id))
        return None
    return scan

def find_resources(object, resource_id):
    resources = []
    if object._urt == 'projects':
        for subject in object.subjects():
            resources.extend(find_resources(subject, resource_id))
    elif object._urt == 'subjects':
        for experiment in object.experiments():
            resources.extend(find_resources(experiment, resource_id))
    elif object._urt == 'experiments':
        for scan in object.scans():
            resources.extend(find_resources(scan, resource_id))
        for assessor in object.assessors():
            for r in assessor.in_resources():
                if r.id() == resource_id or r.label() == resource_id:
                    resources.append(r)
            for r in assessor.out_resources():
                if r.id() == resource_id or r.label() == resource_id:
                    resources.append(r)
        for recon in object.reconstructions():
            for r in recon.in_resources():
                if r.id() == resource_id or r.label() == resource_id:
                    resources.append(r)
            for r in recon.out_resources():
                if r.id() == resource_id or r.label() == resource_id:
                    resources.append(r)
    elif object._urt == 'scans':
        for r in object.resources():
            if r.id() == resource_id or r.label() == resource_id:
                resources.append(r)
    else:
        for r in object.resources():
            if r.id() == resource_id or r.label() == resource_id:
                resources.append(r)
    return resources

def display_parents(object, prefix=''):
    prefix = prefix + '    '
    parent = object.parent()
    type = parent._urt[:-1]
    print '%sfor %s %s' % (prefix, type, parent.label())
    if type != 'project':
        display_parents(parent, prefix)
    return

def get_file(f, path):
    "get a file, but first make directories"
    dirs = []
    dir = os.path.dirname(path)
    while dir:
        dirs.append(dir)
        dir = os.path.dirname(dir)
    dirs.reverse()
    for dir in dirs:
        if not os.path.exists(dir):
            os.mkdir(dir)
    f.get(path)
    return

progname = os.path.basename(sys.argv.pop(0))

try:
    env_user = os.environ['XNAT_USER']
except KeyError:
    env_user = None

try:
    env_password = os.environ['XNAT_PASSWORD']
except KeyError:
    env_password = None

dry_run_flag = False
unauth_flag = False
cl_user = None
cl_password = None
lag = 5

if not sys.argv:
    print
    print 'usage: %s [options] command [command arguments ...]' % progname
    print
    print 'options are:'
    print
    print '    -d -- dry run (no actual changes)'
    print '    -u <user name> -- XNAT user name'
    print '    -p <password> -- XNAT password'
    print '    -n -- use unauthenticated connection'
    print '    -l <time in seconds> -- lag after writing before checking XNAT '
    print '                            modification time (default %d)' % lag
    print
    print 'option values may be given by environment variables:'
    print
    if env_user is None:
        print '    user name: XNAT_USER (not set)'
    else:
        print '    user name: XNAT_USER (set to %s)' % env_user
    if env_password is None:
        print '    password: XNAT_PASSWORD (not set)'
    else:
        print '    password: XNAT_PASSWORD (currently set)'
    print
    print 'if none are given, %s will prompt for them' % progname
    print
    print 'commands and arguments are:'
    print
    print '    list <XNAT URI> <project> [subject [experiment [scan]]] -- list resources'
    print '    init -c <XNAT URI> <resource URI> -- create and associate with a resource'
    print '    init <XNAT URI> <resource URI> -- associate with a resource'
    print '    init <XNAT URI> <project> [subject [experiment [scan]]] <resource ID>'
    print '        -- associate with a resource'
    print '    status -- print file-by-file synchronization status'
    print '    info -- print information about the association'
    print '    push [-f] -- (force) push changes to the server'
    print '    pull [-f] -- (force) pull changes from the server'
    print
    print 'status lines are:'
    print
    print '    local_status remote_status path'
    print
    print 'status codes are:'
    print
    print '    N -- new'
    print '    U -- unchanged'
    print '    M -- modified'
    print '    D -- deleted'
    print '    ? -- modification time is before last sync'
    print '    - -- does not exist (only occurs if the counterpart has status N)'
    print
    print 'push/pull updates are performed as follows (source/dest/action):'
    print
    print '    N N   update if forced'
    print '    N -   update'
    print '    U MD? update if forced'
    print '    M U   update'
    print '    M MD? update if forced'
    print '    D UD  update'
    print '    D M?  update if forced'
    print '    ? UMD update if forced'
    print '    - N   update if forced'
    print
    sys.exit(1)

"""
       DEST
    N U M D ? -
  N f x x x x u
S U x   f f f x
R M x u f f f x
C D x u f u f x
  ? x f f f f x
  - f x x x x x
"""

while sys.argv:
    if not sys.argv[0] or sys.argv[0][0] != '-':
        break
    if sys.argv[0] == '-d':
        dry_run_flag = True
        sys.argv.pop(0)
        continue
    if sys.argv[0] == '-n':
        unauth_flag = True
        sys.argv.pop(0)
        continue
    if sys.argv[0] == '-u':
        sys.argv.pop(0)
        if not sys.argv:
            command_line_error('-u requires an argument')
            sys.exit(1)
        cl_user = sys.argv.pop(0)
        continue
    if sys.argv[0] == '-p':
        sys.argv.pop(0)
        if not sys.argv:
            command_line_error('-p requires an argument')
            sys.exit(1)
        cl_password = sys.argv.pop(0)
        continue
    if sys.argv[0] == '-l':
        sys.argv.pop(0)
        if not sys.argv:
            command_line_error('-l requires an argument')
            sys.exit(1)
        s_lag = sys.argv.pop(0)
        try:
            lag = int(lag)
        except ValueError:
            command_line_error('bad lag value "%s"' % s_lag)
            sys.exit(1)
    command_line_error('unknown option "%s"' % sys.argv[0])
    sys.exit(1)

if not sys.argv:
    command_line_error('no command given')
    sys.exit(1)

if unauth_flag:
    print 'using unauthenticated connection'
    user = None
    password = None
else:
    if cl_user is not None:
        user = cl_user
    elif env_user is not None:
        user = env_user
    else:
        sys.stdout.write('User name: ')
        sys.stdout.flush()
        user = sys.stdin.readline().strip()
    if cl_password is not None:
        password = cl_password
    elif env_password is not None:
        password = env_password
    else:
        password = getpass.getpass()

if not sys.argv:
    command_line_error('no command given')
    sys.exit(1)

command = sys.argv.pop(0)

if command == 'list':
    if not sys.argv:
        command_line_error('no URI given')
        sys.exit(1)
    xnat_uri = sys.argv.pop(0)
    i = connect(xnat_uri, user, password)
    object = get_object(i, sys.argv)
    if object is None:
        sys.exit(1)
    if object._urt == 'projects':
        display_project_resources(object)
    elif object._urt == 'subjects':
        display_subject_resources(object)
    elif object._urt == 'experiments':
        display_experiment_resources(object)
    else:
        display_scan_resources(object)
elif command == 'init':
    if os.path.exists('.xnat'):
        sys.stderr.write('%s: .xnat exists, already a sync dir\n' % progname)
        sys.exit(1)
    if sys.argv[0] == '-c':
        sys.argv.pop(0)
        try:
            xnat_uri = sys.argv.pop(0)
            resource_uri = sys.argv.pop(0)
        except IndexError:
            command_line_error('not enough arguments to init')
            sys.exit(1)
        i = connect(xnat_uri, user, password)
        resource = i.select(resource_uri)
        if resource.exists():
            sys.stderr.write('%s: resource exists\n' % progname)
            sys.exit(1)
        if dry_run_flag:
            print 'not creating resource'
        else:
            print 'creating resource'
            resource.create()
    else:
        xnat_uri = sys.argv.pop(0)
        i = connect(xnat_uri, user, password)
        if len(sys.argv) == 1:
            resource_uri = i.select(resource_uri)
            if not resource.exists():
                sys.stderr.write("%s: can't find resource\n" % progname)
                sys.exit(1)
        else:
            resource_id = sys.argv.pop()
            object = get_object(i, sys.argv)
            print 'searching for resource...'
            resources = find_resources(object, resource_id)
            if not resources:
                sys.stderr.write('%s: resource %s not found\n' % (progname, resource_id))
                sys.exit(1)
            if len(resources) > 1:
                sys.stderr.write('%s: multiple resources %s found\n' % (progname, resource_id))
                sys.exit(1)
            resource = resources[0]
    files = {}
    if dry_run_flag:
        print 'not writing .xnat'
    else:
        print 'writing .xnat'
        write_rc(xnat_uri, resource._uri, files)
elif command == 'status':
    (xnat_uri, resource_uri, last_updated, rc_files) = read_rc()
    i = connect(xnat_uri, user, password)
    resource = i.select(resource_uri)
    if not resource.exists():
        sys.stderr.write("%s: can't find resource\n" % progname)
        sys.exit(1)
    statuses = get_statuses(rc_files, resource)
    for path in sorted(statuses):
        print '%s %s %s' % (statuses[path]['local'], 
                            statuses[path]['remote'], 
                            path)
elif command == 'info':
    (xnat_uri, resource_uri, last_updated, rc_files) = read_rc()
    i = connect(xnat_uri, user, password)
    resource = i.select(resource_uri)
    if not resource.exists():
        sys.stderr.write("%s: can't find resource\n" % progname)
        sys.exit(1)
    print 'Host: %s' % xnat_uri
    print 'Resource URI: %s' % resource_uri
    dt = datetime.datetime.fromtimestamp(last_updated)
    print 'Last synchronized: %s'  % dt.strftime('%a %b %d %H:%M:%S %Y')
    print 'Resource: %s (%s)' % (resource.label(), resource.id())
    display_parents(resource)
elif command == 'push':
    if sys.argv and sys.argv[0] == '-f':
        force_flag = True
    else:
        force_flag = False
    (xnat_uri, resource_uri, last_updated, rc_files) = read_rc()
    i = connect(xnat_uri, user, password)
    resource = i.select(resource_uri)
    if not resource.exists():
        sys.stderr.write("%s: can't find resource\n" % progname)
        sys.exit(1)
    statuses = get_statuses(rc_files, resource)
    updates = []
    for path in sorted(statuses):
        if statuses[path]['local'] == 'N':
            if statuses[path]['remote'] == 'N':
                if not force_flag:
                    print 'conflict: %s' % path
                elif dry_run_flag:
                    print 'not updating %s' % path
                else:
                    print 'updating %s' % path
                    f = resource.file(path)
                    f.put(path)
                    updates.append((path, f))
            if statuses[path]['remote'] == '-':
                if dry_run_flag:
                    print 'not updating %s' % path
                else:
                    print 'updating %s' % path
                    f = resource.file(path)
                    f.put(path)
                    updates.append((path, f))
        if statuses[path]['local'] == 'D':
            if statuses[path]['remote'] == 'D':
                if dry_run_flag:
                    print 'not updating %s' % path
                else:
                    print 'updating %s' % path
                    del rc_files[path]
            if statuses[path]['remote'] == 'M':
                if not force_flag:
                    print 'conflict: %s' % path
                elif dry_run_flag:
                    print 'not updating %s' % path
                else:
                    print 'updating %s' % path
                    f = resource.file(path)
                    f.delete()
                    del rc_files[path]
            if statuses[path]['remote'] == 'U':
                if dry_run_flag:
                    print 'not updating %s' % path
                else:
                    print 'updating %s' % path
                    f = resource.file(path)
                    f.delete()
                    del rc_files[path]
        if statuses[path]['local'] == 'M':
            if statuses[path]['remote'] == 'D':
                if not force_flag:
                    print 'conflict: %s' % path
                elif dry_run_flag:
                    print 'not updating %s' % path
                else:
                    print 'updating %s' % path
                    f = resource.file(path)
                    f.put(path)
                    updates.append((path, f))
            if statuses[path]['remote'] == 'M':
                if not force_flag:
                    print 'conflict: %s' % path
                elif dry_run_flag:
                    print 'not updating %s' % path
                else:
                    print 'updating %s' % path
                    f = resource.file(path)
                    f.put(path)
                    updates.append((path, f))
            if statuses[path]['remote'] == 'U':
                if dry_run_flag:
                    print 'not updating %s' % path
                else:
                    print 'updating %s' % path
                    f = resource.file(path)
                    f.put(path)
                    updates.append((path, f))
        if statuses[path]['local'] == 'U':
            if statuses[path]['remote'] == 'D':
                if not force_flag:
                    print 'conflict: %s' % path
                elif dry_run_flag:
                    print 'not updating %s' % path
                else:
                    print 'updating %s' % path
                    f = resource.file(path)
                    f.put(path)
                    updates.append((path, f))
            if statuses[path]['remote'] == 'M':
                if not force_flag:
                    print 'conflict: %s' % path
                elif dry_run_flag:
                    print 'not updating %s' % path
                else:
                    print 'updating %s' % path
                    f = resource.file(path)
                    f.put(path)
                    updates.append((path, f))
    if not updates:
        print 'nothing updated'
    elif dry_run_flag:
        print 'not updating .xnat'
    else:
        print 'waiting %d seconds...' % lag
        for i in xrange(lag):
            sys.stdout.write('%d...' % (lag-i))
            sys.stdout.flush()
            time.sleep(1)
        print '0'
        for (path, f) in updates:
            print 'getting times for %s' % path
            rc_files[path] = files_times(path, f)
        write_rc(xnat_uri, resource_uri, rc_files)
elif command == 'pull':
    if sys.argv and sys.argv[0] == '-f':
        force_flag = True
    else:
        force_flag = False
    (xnat_uri, resource_uri, last_updated, rc_files) = read_rc()
    i = connect(xnat_uri, user, password)
    resource = i.select(resource_uri)
    if not resource.exists():
        sys.stderr.write("%s: can't find resource\n" % progname)
        sys.exit(1)
    statuses = get_statuses(rc_files, resource)
    prunes = []
    for path in sorted(statuses):
        if statuses[path]['local'] == '?' or statuses[path]['remote'] == '?':
            print 'panic: %s' % path
            prunes.append(path)
        elif statuses[path]['local'] == 'N' and statuses[path]['remote'] != '-':
            print 'panic: %s' % path
            prunes.append(path)
        elif statuses[path]['remote'] == 'N' and statuses[path]['local'] != '-':
            print 'panic: %s' % path
            prunes.append(path)
        elif statuses[path]['local'] == '-' and statuses[path]['remote'] != 'N':
            print 'panic: %s' % path
            prunes.append(path)
        elif statuses[path]['remote'] == '-' and statuses[path]['local'] != 'N':
            print 'panic: %s' % path
            prunes.append(path)
    for path in prunes:
        del statuses[path]
    for path in sorted(statuses):
        if statuses[path]['remote'] == 'N':
            if statuses[path]['local'] == 'N':
                if not force_flag:
                    print 'conflict: %s' % path
                elif dry_run_flag:
                    print 'not updating %s' % path
                else:
                    print 'updating %s' % path
                    f = resource.file(path)
                    get_file(f, path)
                    rc_files[path] = files_times(path, f)
            if statuses[path]['local'] == '-':
                if dry_run_flag:
                    print 'not updating %s' % path
                else:
                    print 'updating %s' % path
                    f = resource.file(path)
                    get_file(f, path)
                    rc_files[path] = files_times(path, f)
        if statuses[path]['remote'] == 'D':
            if statuses[path]['local'] == 'D':
                if dry_run_flag:
                    print 'not updating %s' % path
                else:
                    print 'updating %s' % path
                    del rc_files[path]
            if statuses[path]['local'] == 'M':
                if not force_flag:
                    print 'conflict: %s' % path
                elif dry_run_flag:
                    print 'not updating %s' % path
                else:
                    print 'updating %s' % path
                    os.unlink(path)
                    del rc_files[path]
            if statuses[path]['local'] == 'U':
                if dry_run_flag:
                    print 'not updating %s' % path
                else:
                    print 'updating %s' % path
                    os.unlink(path)
                    del rc_files[path]
        if statuses[path]['remote'] == 'M':
            if statuses[path]['local'] == 'D':
                if not force_flag:
                    print 'conflict: %s' % path
                elif dry_run_flag:
                    print 'not updating %s' % path
                else:
                    print 'updating %s' % path
                    f = resource.file(path)
                    get_file(f, path)
                    rc_files[path] = files_times(path, f)
            if statuses[path]['local'] == 'M':
                if not force_flag:
                    print 'conflict: %s' % path
                elif dry_run_flag:
                    print 'not updating %s' % path
                else:
                    print 'updating %s' % path
                    f = resource.file(path)
                    get_file(f, path)
                    rc_files[path] = files_times(path, f)
            if statuses[path]['local'] == 'U':
                if dry_run_flag:
                    print 'not updating %s' % path
                else:
                    print 'updating %s' % path
                    f = resource.file(path)
                    get_file(f, path)
                    rc_files[path] = files_times(path, f)
        if statuses[path]['remote'] == 'U':
            if statuses[path]['local'] == 'D':
                if not force_flag:
                    print 'conflict: %s' % path
                elif dry_run_flag:
                    print 'not updating %s' % path
                else:
                    print 'updating %s' % path
                    f = resource.file(path)
                    get_file(f, path)
                    rc_files[path] = files_times(path, f)
            if statuses[path]['local'] == 'M':
                if not force_flag:
                    print 'conflict: %s' % path
                elif dry_run_flag:
                    print 'not updating %s' % path
                else:
                    print 'updating %s' % path
                    f = resource.file(path)
                    get_file(f, path)
                    rc_files[path] = files_times(path, f)
    if dry_run_flag:
        print 'not writing .xnat'
    else:
        print 'writing .xnat'
        write_rc(xnat_uri, resource_uri, rc_files)
else:
    command_line_error('unknown command "%s"' % command)
    sys.exit(1)

sys.exit(0)

# eof