SlicerLabkeyExtension/labkeyBrowser/slicerNetwork.py

import qt
import urllib2
import socket
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 ConnectionError(Exception):
    def __init__(self,message):
        self.message=message


class slicerNetwork(slicer.ScriptedLoadableModule.ScriptedLoadableModule):
    def __init__(self, parent):
        slicer.ScriptedLoadableModule.ScriptedLoadableModule.__init__(self,parent)
        self.parent.title="slicerNetwork"
        pass

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

        self.localCacheDirectory=os.path.join(fhome,"labkeyCache")
        self.configDir=os.path.join(fhome,".labkey")
        self.mode="http"
        #try initializing network from default config file, if found
        self.initFromConfig()


    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"

    #configuration part

    def configureSSL(self,cert,key,pwd,cacert):
    #do this first
        try:
            self.ctx=ssl.SSLContext(ssl.PROTOCOL_SSLv23)
            self.ctx.load_cert_chain(cert,keyfile=key,password=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 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 initFromConfig(self):
        path=os.path.join(self.configDir,"Remote.json")
        try:
            self.parseConfig(path)
        except OSError:
            return
        self.initRemote()

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

        dt=json.load(f)
        self.mode="http"
        if dt.has_key('SSL'):
            self.configureSSL(
                dt['SSL']['user'],
                dt['SSL']['key'],
                dt['SSL']['keyPwd'],
                dt['SSL']['ca']
                )
            if not parent==None:
                parent.userCertButton.setText(dt['SSL']['user'])
                parent.caCertButton.setText(dt['SSL']['ca'])
                parent.privateKeyButton.setText(dt['SSL']['key'])
                parent.pwd=dt['SSL']['keyPwd']


        self.hostname=dt['host']
        self.auth_name=dt['labkey']['user']
        self.auth_pass=dt['labkey']['password']
        if not parent==None:
            parent.serverURL.setText(dt['host'])
            parent.authName.setText(dt['labkey']['user'])
            parent.authPass.setText(dt['labkey']['password'])


    #path convention

    #localPath is a path on local filesystem. It means it has to adhere to OS
    #policy, especially in separators

    #relativePath is just the portion of the path that is equal both locally
    #and remotely. In relativePath, separator is according to http convention,
    #which equals separator in osx and *nix

    #labkeyPath or remotePath is the full URL to the resource,
    #including http-like header and all intermediate portions

    #the following functions convert between different notations

    #was GetLocalPath
    def GetLocalPathFromRelativePath(self,relativePath):
        debug=False
        relativePath=re.sub('labkey://','',relativePath)
        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.GetLocalCacheDirectory(),relativePath)

    #was GetRemotePath
    def GetLabkeyPathFromRelativePath(self,source):
        return self.GetLabkeyWebdavUrl()+"/"+source

    #was GetLabkeyPathFromLocalPath
    def GetRelativePathFromLocalPath(self,f):
        debug=False
        #report it with URL separator, forward-slash
        if f.find(self.localCacheDirectory)<-1:
            print("Localpath misformation. Exiting")
            return "NULL"
        f0=re.sub('\\\\','\\\\\\\\',self.localCacheDirectory)
        relativePath=re.sub(re.compile(f0),'',f)
        if debug:
            print("[SEP] Relative path {}").format(relativePath)
        #leaves /ContextPath/%40files/subdirectory_list/files
        #remove first separator
        sp=os.path.sep
        f1=re.sub('\\\\','\\\\\\\\',sp)
        if relativePath[0]==sp:
            relativePath=relativePath[1:len(relativePath)]

        if debug:
            print("[SLASH] Relative path {}").format(relativePath)

        return re.sub(f1,'/',relativePath)

    #was GetLabkeyPathFromRemotePath
    def GetRelativePathFromLabkeyPath(self,f):
        #used to query labkey stuff, so URL separator is used

        #in old conventions, labkey://was used to tag a labkeyPath
        f=re.sub('labkey://','',f)
        f=re.sub(self.GetLabkeyWebdavUrl(),'',f)

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


    #standard HTTP
    def get(self,url):

        debug=False
        if debug:
            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)
        try:
            return self.opener.open(r)
        #f contains json as a return value
        except urllib2.HTTPError as e:
            print e.code
            print e.read()
            return e
        except urllib2.URLError as e:
            print("Handling URLError")
            print e
            return None
    
    #a HEAD request
    def head(self,url):
    
        debug=False
        if debug:
            print("HEAD: {0}").format(url)
            print("as {0}").format(self.auth_name)
        r=MethodRequest(url.encode('utf-8'),method="HEAD")
        base64string = base64.b64encode('%s:%s' % (self.auth_name, self.auth_pass))
        r.add_header("Authorization", "Basic %s" % base64string)
        try:
            return self.opener.open(r)
        #f contains json as a return value
        except urllib2.HTTPError as e:
            print e.code
            print e.read()
            return e


    def post(self,url,data):

        debug=False
        r=urllib2.Request(url)
        #makes it a post
        r.add_data(data)
        r.add_header("Content-Type","application/json")
        #add csrf
        csrf=self.getCSRF()
        r.add_header("X-LABKEY-CSRF",csrf)

        base64string = base64.b64encode('%s:%s' % (self.auth_name, self.auth_pass))
        r.add_header("Authorization", "Basic %s" % base64string)
        if debug:
            print("{}: {}").format(r.get_method(),r.get_full_url())
            print("data: {}").format(r.get_data())
            print("Content-Type: {}").format(r.get_header('Content-Type'))
        try:
            return self.opener.open(r)
        except urllib2.HTTPError as e:
            print e.code
            print e.read()
            return e
        #f contains json as a return value

    def put(self,url,data):

        debug=False

        if debug:
            print("PUT: {}").format(url)
        r=MethodRequest(url.encode('utf-8'),method="PUT")
        #makes it a post
        r.add_data(data)
        print("PUT: data size: {}").format(len(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 getCSRF(self):
        url=self.GetLabkeyUrl()+'/login/whoAmI.view'
        try:
            jsonData=json.load(self.get(url))
        except AttributeError:
            print("Failed")
            return None
        
        return jsonData["CSRF"]


    #file manipulation routiens

    #was GetFile
    def DownloadFileToCache(self,relativePath):
        debug=False
        # check for file in cache. If available, use, if not, download
        localPath=self.GetLocalPathFromRelativePath(relativePath)
        if os.path.isfile(localPath):
            if debug:
                print("Returning localFile {}".format(localPath))
            return localPath

        if debug:
            print("labkeyURIHandler::DownloadFileToCache({0}->{1})").format(relativePath,localPath)

        #make all necessary directories LOCALLY
        path=os.path.dirname(localPath)
        if debug:
            print("Target directory: '{}''".format(path))
        if not os.path.isdir(path):
            os.makedirs(path)

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


        #labkeyPath=self.GetLabkeyPathFromRelativePath(relativePath)
        remoteBuffer=self.readFileToBuffer(relativePath)

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


        remoteBuffer.seek(0)

        shutil.copyfileobj(remoteBuffer,localBuffer)
        if debug:
	         print("Local size: {0}").format(localBuffer.tell())
        localBuffer.close()
        return localPath

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

    #mimic slicer.util.loadNodeFromFile
    def loadNode(self, relativeName, filetype, properties={},returnNode=False, keepCached=True):
         #this is the only relevant part - file must be downloaded to cache
         #labkeyName is just the relative part (from labkey onwards)
         
         localPath=self.DownloadFileToCache(relativeName)
         
         if not returnNode:
             slicer.util.loadNodeFromFile(localPath,filetype,properties,returnNode=False)
             if not keepCached:
                 os.remove(localPath)
             return
         
         ok,node=slicer.util.loadNodeFromFile(localPath,filetype,properties,returnNode=True)
         if ok:
             if not keepCached:
                 os.remove(localPath)
             return node
         

         os.remove(localPath)
         return None
    #     #remove retrieved file

    def storeNode(self,node,project,dir=None):

        removeFile=True
        relativePath=project+'/%40files'
        if not dir==None:
            relativePath+='/'+dir
        labkeyPath=self.GetLabkeyPathFromRelativePath(relativePath)
        print ("Remote: {}").format(labkeyPath)
      
        #checks if exists
        self.mkdir(labkeyPath)

        localPath=self.GetLocalPathFromRelativePath(relativePath)
        localPath.replace('/',os.path.sep)

        nodeName=node.GetName()

        suffix=".nrrd"
        if node.__class__.__name__=="vtkMRMLDoubleArrayNode":
            suffix=".mcsv"
        if node.__class__.__name__=="vtkMRMLTransformNode":
            suffix=".h5"
        fileName=nodeName+suffix
        file=os.path.join(localPath,fileName)
        slicer.util.saveNode(node,file)
        print("Stored to: {}").format(file)
        f=open(file,"rb")

        f.seek(0,2)
        sz=f.tell()
        print("File size in memory: {0}").format(sz)
        f.seek(0)

        localBuffer=f.read()
        print("Local buffer size: {}").format(len(localBuffer))
        remoteFile=labkeyPath+'/'+fileName
        resp=self.put(remoteFile,localBuffer)
        print(resp.read())
        f.close()
        if removeFile:
            os.remove(file)



    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):
        relativePath=self.GetRelativePathFromLabkeyPath(remoteDir)
        s=0
        while True:
            s1=relativePath.find('/',s)
            if s1<0:
                break
            path=relativePath[0:s1]
            remotePath=self.GetLabkeyPathFromRelativePath(relativePath)
            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

    #was listDir
    def listRelativeDir(self,relativeDir):
        print("Listing for {0}").format(relativeDir)
        dirUrl=self.GetLabkeyPathFromRelativePath(relativeDir)
        status,dirs=self.listRemoteDir(dirUrl)
        dirs=[self.GetRelativePathFromLabkeyPath(d) for d in dirs];
        return dirs

    #was isDir
    def isRemoteDir(self, remotePath):
        #print "isDir: {}".format(remotePath)
        r=MethodRequest(remotePath,method="PROPFIND")
        PROPFIND=u"""<?xml version="1.0" encoding="utf-8"?>\n
            <a:propfind xmlns:a="DAV:">\n
            <a:prop>\n
            <a:resourcetype/>\n
            </a:prop>\n
            </a:propfind>"""
        r.add_header('content-type','text/xml; charset="utf-8"')
        r.add_header('content-length',str(len(PROPFIND)))
        r.add_header('Depth','0')
        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(remotePath)
        try:
            f=self.opener.open(r)
        except:
            return False
        tree=ET.XML(f.read())
        try:
            rps=tree.find('{DAV:}response').find('{DAV:}propstat')
            rps=rps.find('{DAV:}prop')
            rps=rps.find('{DAV:}resourcetype').find('{DAV:}collection')
            if rps != None:
                return True
            else:
                return False
        except:
            return False

    def listRemoteDir(self,dirUrl):
        #input is remoteDir, result are remoteDirs

        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)
            dirent=self.GetHostName()+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 readFileToBuffer(self, relativePath):
        dirUrl=self.GetLabkeyPathFromRelativePath(relativePath)
        f=self.get(dirUrl)
        return StringIO.StringIO(f.read())

    def uploadFile(self,localPath):
        #get upstream directories sorted out
        relativePath=self.GetRelativePathFromLocalPath(localPath)
        if relativePath=="NULL":
            errorCode="Failed to upload {}. Potential incorrect location"
            errorCode+=". Should be in labkeyCache!"
            print(errorCode.format(relativePath))
            return False
        relativeDir=relativePath[0:labkeyPath.rfind('/')]
        remoteDir=self.GetLabkeyPathFromRemotePath(relativeDir)
        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, 'rb') as f:
            data=f.read()
        remotePath=self.GetLabkeyPathFromRelativePath(relativePath)
        self.put(remotePath,data)

    #was copyFileToRemote
    def copyLocalFileToRemote(self,localPath, remotePath):
        #get upstream directories sorted out

        remoteDir=os.path.dirname(remotePath)
        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
        if (self.isRemoteDir(remotePath)):
            remotePath=remotePath+'/'+os.path.basename(localPath)

        with open(localPath, 'rb') as f:
            data=f.read()

        itry=0
        while itry<10:
            try:
                self.put(remotePath,data)
                break
            except (socket.error,urllib2.URLError):
                print("[{}] Failed to execute put(), retry".format(itry))
                itry=itry+1
        if itry==10:
            raise ConnectionError('Could not copy file, number of tries exceeded') 

    #was loadDir
    def DownloadDirToCache(self, relativePath):

        files=self.listRelativeDir(relativePath)
        fdir="NONE"
        for f in files:
            #f is local path
            try:
                localDir=os.path.dirname(self.DownloadFileToCache(f))
            except:
                #fails if there is a subdirectory; go recursively
                print("self.readDir(f) not implemented")
        return localDir

    #database routines

    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 filterDataset(self,project,dataset,filter):
        debug=False
        url=self.GetLabkeyUrl()+'/'+project
        url+='/query-selectRows.api?schemaName=study&query.queryName='+dataset
        for f in filter:
            url+="&query."+f['variable']+"~"+f['oper']+"="+f['value']
        if debug:
            print("Sending {}").format(url)
        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))

    def filterList(self,project,dataset,filter):
        schemaName='lists'
        debug=False
        url=self.GetLabkeyUrl()+'/'+project
        url+='/query-selectRows.api?schemaName='+schemaName+'&query.queryName='+dataset
        for f in filter:
            url+="&query."+f['variable']+"~"+f['oper']+"="+f['value']
        if debug:
            print("Sending {}").format(url)
        return json.load(self.get(url))


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