SlicerLabkeyExtension/labkeySlicerPythonExtension/slicerNetwork.py

import qt
import urllib2
import ssl
import cookielib
import xml.etree.ElementTree as ET
import re
import StringIO
import slicer
import shutil
import distutils
import os
import base64
import json

#see https://gist.github.com/logic/2715756, allow requests to do Put
class MethodRequest(urllib2.Request):
    def __init__(self, *args, **kwargs):
        if 'method' in kwargs:
            self._method = kwargs['method']
            del kwargs['method']
        else:
            self._method = None
        return urllib2.Request.__init__(self, *args, **kwargs)

    def get_method(self, *args, **kwargs):
        if self._method is not None:
            return self._method
        return urllib2.Request.get_method(self, *args, **kwargs)

class slicerNetwork:
    def __init__(self, *args, **kwargs):
        pass

class labkeyURIHandler(slicer.vtkURIHandler):
    def __init__(self):
        slicer.vtkURIHandler.__init__(self)
        self.className="labkeyURIHandler"
        slicer.mrmlScene.AddURIHandler(self)
        try:
	       self.localCacheDirectory=os.path.join(os.environ["HOME"],"labkeyCache")
        except:
	    #in windows, the variable is called HOMEPATH
	       fhome=os.environ['HOMEDRIVE']+os.environ['HOMEPATH']
	       self.localCacheDirectory=os.path.join(fhome,"labkeyCache")
        self.mode="http"

    def CanHandleURI(self,uri):
        print "labkeyURIHandler::CanHandleURI({0})".format(uri)
        if uri.find('labkey://')==0:
            return 1
        return 0

    def GetClassName(self):
        return self.className

    def GetHostName(self):
        return self.hostname

    def SetHostName(self,host):
        self.hostname=host

    def SetLocalCahceDirectory(self,dir):
        self.localCacheDirectory=dir

    def GetLocalCacheDirectory(self):
        return self.localCacheDirectory

    def GetLabkeyUrl(self):
        return self.hostname+"/labkey"

    def GetLabkeyWebdavUrl(self):
        return self.GetLabkeyUrl()+"/_webdav"

    def GetLocalPath(self,source):
        debug=False
        relativePath=re.sub('labkey://','',source)
        sp=os.sep.encode('string-escape')
        if debug:
            print "Substituting / with {0} in {1}".format(sp,relativePath)
        relativePath=re.sub('/',sp,relativePath)
        return os.path.join(self.localCacheDirectory,relativePath)

    def GetRemotePath(self,source):
        return self.GetLabkeyWebdavUrl()+"/"+GetLabkeyPathFromLocalPath(source)

    def GetLabkeyPathFromLocalPath(self,f):
        #report it with URL separator, forward-slash
        if f.find(self.localCacheDirectory)<-1:
            print "Localpath misformation. Exiting"
            return "NULL"
        dest=re.sub(self.localCacheDirectory,'',f)
        #leaves /ContextPath/%40files/subdirectory_list/files
        #remove first separator
        sp=os.sep.encode('string-escape')
        if dest[0]==sp:
            dest=dest[1:len(dest)]
        return re.sub(sp,'/',dest)

    def GetLabkeyPathFromRemotePath(self,f):
        #used to query labkey stuff, so URL separator is used
        f=re.sub('labkey://','',f)
        f=re.sub(self.GetLabkeyWebdavUrl(),'',f)

        if f[0]=='/':
            f=f[1:len(f)]
        return f;


    def GetFile(self,source):
        # check for file in cache. If available, use, if not, download
        localPath=self.GetLocalPath(source)
        if os.path.isfile(localPath):
            return localPath
        self.StageFileRead(source,localPath)
        return localPath

    def StageFileRead(self,source,dest):
        print "labkeyURIHandler::StageFileRead({0},{1})".format(source,dest)
        labkeyPath=re.sub('labkey://','',source)
        remote=self.readFile(self.hostname,labkeyPath)
        #make all necessary directories
        path=os.path.dirname(dest)
        if not os.path.isdir(path):
            os.makedirs(path)

        local=open(dest,'wb')
        #make sure we are at the begining of the file

	#check file size
	remote.seek(0,2)
	sz=remote.tell()
	print "Remote size: {0}".format(sz)

	remote.seek(0)
        shutil.copyfileobj(remote,local)
	print "Local size: {0}".format(local.tell())
        local.close()

    def StageFileWrite(self,source,dest):
        print "labkeyURIHandler::StageFileWrite({0},{1}) not implemented yet".format(
            source,dest)

    def fileTypesAvailable(self):
        return ('VolumeFile','SegmentationFile','TransformFile')

    #mimic slicer.util.loadNodeFromFile
    def loadNodeFromFile(self, filename, filetype, properties={}, returnNode=False):
        #this is the only relevant part - file must be downloaded to cache
        localPath=self.GetFile(filename)
        slicer.util.loadNodeFromFile(localPath,filetype,properties,returnNode)
        #remove retrieved file
        try:
            if not(properties['keepCachedFile']) :
                os.remove(localPath)
        except:
            pass


    def loadVolume(self,filename, properties={}, returnNode=False):
        filetype = 'VolumeFile'
        #redirect to self.loadNodeFromFile first to get the cached file
        return self.loadNodeFromFile(filename,filetype, properties,returnNode)

    def loadSegmentation(self,filename,properties={},returnNode=False):
        filetype='SegmentationFile'
        #redirect to self.loadNodeFromFile first to get the cached file
        return self.loadNodeFromFile(filename,filetype, properties,returnNode)
    #add others if needed

    ## setup & initialization routines

    def configureSSL(self,cert,key,pwd,cacert):
        #do this first
        try:
            self.ctx=ssl.SSLContext(ssl.PROTOCOL_SSLv23)
            self.ctx.load_cert_chain(cert,key,pwd)
            self.ctx.verify_mode=ssl.CERT_REQUIRED
            self.ctx.load_verify_locations(cacert)
        except ssl.SSLError as err:
            print " Failed to configure SSL: {0}".format(str(err))
        self.mode="https"


    def connectRemote(self,serverUrl,uname,pwd):
        https_handler=urllib2.HTTPSHandler(context=self.ctx)
        self.SetHostName(serverUrl)
        #cookie part
        cj=cookielib.CookieJar()
        cookie_handler=urllib2.HTTPCookieProcessor(cj)
        self.opener=urllib2.build_opener(https_handler,cookie_handler)
        loginUrl=serverUrl+"/labkey/login/login.post"
        r=urllib2.Request(str(loginUrl))
        print "Connecting to {0}".format(loginUrl)
        r.add_header('ContentType','application/x-www-form-urlencoded')
        data="email="+uname+"&password="+pwd
        #print "Data: '{0}".format(data)
        r.add_data(data)

        self.opener.open(r)

    def initRemote(self):
        if self.mode=="https":
            http_handler=urllib2.HTTPSHandler(context=self.ctx)
        if self.mode=="http":
            http_handler=urllib2.HTTPHandler()
        #cookie part
        cj=cookielib.CookieJar()
        cookie_handler=urllib2.HTTPCookieProcessor(cj)
        self.opener=urllib2.build_opener(http_handler,cookie_handler)

    def parseConfig(self,fname):
        try:
            f=open(fname)
        except OSError as e:
            print "Confgiuration error: OS error({0}): {1}".format(e.errno, e.strerror)
            return

        dt=json.load(f)
        if dt.has_key('SSL'):
            self.configureSSL(
                dt['SSL']['user'],
                dt['SSL']['key'],
                dt['SSL']['keyPwd'],
                dt['SSL']['ca']
                )
        self.hostname=dt['host']
        self.auth_name=dt['labkey']['user']
        self.auth_pass=dt['labkey']['password']



    #cj in opener contains the cookies

    def get(self,url):

        print "GET: {0}".format(url)
        print "as {0}".format(self.auth_name)
        r=urllib2.Request(url)
        base64string = base64.b64encode('%s:%s' % (self.auth_name, self.auth_pass))
        r.add_header("Authorization", "Basic %s" % base64string)
        return self.opener.open(r)
        #f contains json as a return value

    def post(self,url,data):

        r=urllib2.Request(url)
        #makes it a post
        r.add_data(data)
        r.add_header("Content-Type","application/json")
        base64string = base64.b64encode('%s:%s' % (self.auth_name, self.auth_pass))
        r.add_header("Authorization", "Basic %s" % base64string)
        print "{}: {}".format(r.get_method(),r.get_full_url())
        print "data: {}".format(r.get_data())
        print "Content-Type: {}".format(r.get_header('Content-Type'))
        return self.opener.open(r)
        #f contains json as a return value

    def put(self,url,data):

        print "PUT: {0}".format(url)
        r=MethodRequest(url,method="PUT")
        #makes it a post
        r.add_data(data)
        r.add_header("content-type","application/octet-stream")
        base64string = base64.b64encode('%s:%s' % (self.auth_name, self.auth_pass))
        r.add_header("Authorization", "Basic %s" % base64string)

        return self.opener.open(r)
        #f contains json as a return value

    def remoteDirExists(self,url):
        status,dirs=self.listRemoteDir(url);
        return status

    def mkdir(self,remoteDir):
        if self.remoteDirExists(remoteDir):
            return False
        r=MethodRequest(remoteDir,method="MKCOL")
        base64string = base64.b64encode('%s:%s' % (self.auth_name, self.auth_pass))
        r.add_header("Authorization", "Basic %s" % base64string)
        try:
            f=self.opener.open(r)
        except:
            print "Error: Failed MKCOL {}".format(remoteDir)
            return False
        return True

    def mkdirs(self,remoteDir):
        labkeyPath=self.GetLabkeyPathFromRemotePath(remoteDir)
        s=0
        while True:
            s1=labkeyPath.find('/',s)
            if s1<0:
                break
            path=labkeyPath[0:s1]
            remotePath=self.GetLabkeyWebdavUrl()+'/'+path
            dirExists=self.remoteDirExists(remotePath)
            if not dirExists:
                if not self.mkdir(remotePath):
                    return False
            s=s1+1
        return self.mkdir(remoteDir)

    def rmdir(self,remoteDir):
        if not self.remoteDirExists(remoteDir):
            return True
        r=MethodRequest(remoteDir,method="DELETE")
        base64string = base64.b64encode('%s:%s' % (self.auth_name, self.auth_pass))
        r.add_header("Authorization", "Basic %s" % base64string)
        try:
            f=self.opener.open(r)
        except:
            print "Error: Failed DELETE {}".format(remoteDir)
            return False
        return True

    def listDir(self,dir):
        print "Listing for {0}".format(dir)
        dirUrl=self.GetLabkeyWebdavUrl()+"/"+dir
        status,dirs=self.listRemoteDir(dirUrl)
        return dirs

    def listRemoteDir(self,dirUrl):
        r=MethodRequest(dirUrl,method="PROPFIND")
        PROPFIND=u"""<?xml version="1.0" encoding="utf-8"?>\n
                    <propfind xmlns="DAV:">\n
                    <prop>\n
                    <getetag/>\n
                    </prop>\n
                    </propfind>"""
        r.add_header('content-type','text/xml; charset="utf-8"')
        r.add_header('content-length',str(len(PROPFIND)))
        r.add_header('Depth','1')
        r.add_data(PROPFIND)
        base64string = base64.b64encode('%s:%s' % (self.auth_name, self.auth_pass))
        r.add_header("Authorization", "Basic %s" % base64string)
        print "PROPFIND: {0}".format(dirUrl)
        dirs=[]
        try:
            f=self.opener.open(r)
        except:
            return False,dirs
        tree=ET.XML(f.read())
        rps=tree.findall('{DAV:}response')
        for r in rps:
            hr=r.find('{DAV:}href')
            dirent=hr.text
            dirent=re.sub('/labkey/_webdav/','',dirent)
            dirs.append(dirent)
        del dirs[0]
        return True,dirs

    def toRelativePath(self,dirs):
        flist1=[]
        for d in dirs:
            if d[-1]=='/':
                d=d[:-1]
            if d.rfind('/')>-1:
                d=d[d.rfind('/')+1:]
            flist1.append(d)
        return flist1

    def readFile(self, serverUrl, path):
        dirUrl=serverUrl+"/labkey/_webdav"+"/"+path
        f=self.get(dirUrl)
        return StringIO.StringIO(f.read())

    def uploadFile(self,localPath):
        #get upstream directories sorted out
        labkeyPath=self.GetLabkeyPathFromLocalPath(localPath)
        if labkeyPath=="NULL":
            errorCode="Failed to upload {}. Potential incorrect location"
            errorCode+=". Should be in labkeyCache!"
            print errorCode.format(localPath)
            return False
        labkeyDir=labkeyPath[0:labkeyPath.rfind('/')]
        remoteDir=self.GetLabkeyWebdavUrl()+'/'+labkeyDir
        if not self.remoteDirExists(remoteDir):
            if not self.mkdirs(remoteDir):
                errorCode="UploadFile: Could not create directory {}"
                print errorCode.format(remoteDir)
                return False

        #make an URL request
        with open(localPath, 'r') as f:
            data=f.read()
        remotePath=self.GetLabkeyWebdavUrl()+'/'+labkeyPath
        self.put(remotePath,data)

    def loadDir(self, path):
        #dirURL=serverUrl+"/labkey/_webdav/"+path
        files=self.listDir(path)
        fdir="NONE"
        for f in files:
            #returns local path
            try:
                fdir=os.path.dirname(self.GetFile(f))
            except:
                #fails if there is a subdirectory; go recursively
                self.readDir(f)
        return fdir

    def loadDataset(self,project,dataset):
        url=self.GetLabkeyUrl()+'/'+project
        url+='/query-selectRows.api?schemaName=study&query.queryName='+dataset
        return json.load(self.get(url))

    def modifyDataset(self,method,project,dataset,rows):
        #method can be insert or update
        data={}
        data['schemaName']='study'
        data['queryName']=dataset
        data['rows']=rows
        url=self.GetLabkeyUrl()+'/'+project
        url+='/query-'+method+'Rows.api?'
        return self.post(url,json.dumps(data))