SlicerLabkeyExtension/labkeyBrowser/slicerNetwork.py

import qt
import urllib3
import xml.etree.ElementTree as ET
import re
import slicer
import shutil
import distutils
import os
import json
import chardet
import io

#mimics labkeyInterface
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)

        fhome=os.path.expanduser('~')
        self.localCacheDirectory=os.path.join(fhome,"labkeyCache")
        self.configDir=os.path.join(fhome,".labkey")
        #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):
        self.hostname=self.connectionConfig['host']
        print("Host name:{}".format(self.hostname))
        return self.hostname

    def SetHostName(self,host):
        try:
            self.connectionConfig['host']=host
        except AttributeError:
            self.connectionConfig={}
            self.connectionConfig['host']=host

        self.hostname=host

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

    def GetLocalCacheDirectory(self):
        return self.localCacheDirectory

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

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

    #configuration part

    def configureSSL(self,cert=None,key=None,pwd=None,cacert=None):
        print("configure SSL: {}".format(cert))

        if not cert==None:
            self.http = urllib3.PoolManager(\
                    cert_file=cert,\
                    cert_reqs='CERT_REQUIRED',\
                    key_file=key,\
                    ca_certs=cacert)
        else:
            self.http = urllib3.PoolManager(\
                    cert_reqs='CERT_REQUIRED',\
                    ca_certs=cacert)

    def initRemote(self):
        #assume self.connectionConfig is set
        self.http=urllib3.PoolManager()
        try:
            config=self.connectionConfig
        except AttributeError:
            print("ConnectionConfig not initialized")
            return
        
        if 'SSL' in self.connectionConfig:
            print("Setting up SSL")
            try:
                cert=self.connectionConfig['SSL']['user']
            except KeyError:
                print("No user cert supplied")
                return
            try:
                key=self.connectionConfig['SSL']['key']
            except KeyError:
                print("No user key supplied")
                return
            try:
                ca=self.connectionConfig['SSL']['ca']
            except KeyError:
                print("No CA cert supplied")
                return
                
            self.configureSSL(cert,key,None,ca)


    def initFromConfig(self):
        print("slicerNetwork: initFromConfig")
        path=os.path.join(self.configDir,"Remote.json")
        try:
            self.parseConfig(path)
        except OSError:
            return
        self.initRemote()

    def parseConfig(self,fname,parent=None):
        print("slicerNetwork: parseConfig") 
        try:
            with open(fname) as f:
                self.connectionConfig=json.load(f)
        except OSError as e:
            print("Confgiuration error: OS error({0}): {1}".format(e.errno, e.strerror))
            raise
        except AttributeError:
            print("Troubles parsing json at {}".format(fname))

        if not parent==None:
            parent.userCertButton.setText(self.connectionConfig['SSL']['user'])
            parent.caCertButton.setText(self.connectionConfig['SSL']['ca'])
            parent.privateKeyButton.setText(self.connectionConfig['SSL']['key'])
            parent.pwd=self.connectionConfig['SSL']['keyPwd']


        self.hostname=self.connectionConfig['host']
        self.auth_name=self.connectionConfig['labkey']['user']
        self.auth_pass=self.connectionConfig['labkey']['password']
        if not parent==None:
            parent.serverURL.setText(self.connectionConfig['host'])
            parent.authName.setText(self.connectionConfig['labkey']['user'])
            parent.authPass.setText(self.connectionConfig['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('unicode_escape').decode('utf-8')
        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;
    
    def getBasicAuth(self):

        user=self.connectionConfig['labkey']['user']
        pwd=self.connectionConfig['labkey']['password']
        #debug
        return user+":"+pwd
     

    #standard HTTP
    def get(self,url,binary=False):

        debug=True
        if debug:
            print("GET: {0}".format(url))
            print("as {0}".format(self.auth_name))
        headers=urllib3.util.make_headers(basic_auth=self.getBasicAuth())
        try:
            if not binary:
                return self.http.request('GET',url,headers=headers)
            else:
                return self.http.request('GET',url,headers=headers,preload_content=False)
        #f contains json as a return value
        #f contains json as a return value
        except urllib3.exceptions.HTTPError as e:
            print(e)
            return e

    #a HEAD request
    def head(self,url):
    
        debug=False
        if debug:
            print("HEAD: {0}".format(url))
            print("as {0}".format(self.auth_name))
        headers=urllib3.util.make_headers(basic_auth=self.getBasicAuth())
        
        try:
            return self.http.request('HEAD',url,headers=headers)
        except urllib3.exceptions.HTTPError as e:
            print(e)
            return e



    def post(self,url,data):

        debug=False
        headers=urllib3.util.make_headers(basic_auth=self.getBasicAuth())
        headers["Content-Type"]="application/json"
        #add csrf;also sets self.cookie
        headers["X-LABKEY-CSRF"]=self.getCSRF()
        headers["Cookie"]=self.cookie

        try:
            return self.http.request('POST',url,headers=headers,body=data)
        #f contains json as a return value
        except urllib3.exceptions.HTTPError as e:
            print(e)
            return e


    def put(self,url,data):
   
        debug=False
 
        if debug:
            print("PUT: {}".format(url))

        headers=urllib3.util.make_headers(basic_auth=self.getBasicAuth())
        headers["Content-Type"]="application/octet-stream"
        #add csrf
        headers["X-LABKEY-CSRF"]=self.getCSRF()
        headers["Cookie"]=self.cookie

        try:
            return self.http.request('PUT',url,headers=headers,body=data)
        #f contains json as a return value
        except urllib3.exceptions.HTTPError as e:
            print(e)
            return e
    
    def sendRequest(self,url,requestCode):

        debug=False
        headers=urllib3.util.make_headers(basic_auth=self.getBasicAuth())
        #add csrf;also sets self.cookie
        headers["X-LABKEY-CSRF"]=self.getCSRF()
        headers["Cookie"]=self.cookie

        try:
            return self.http.request(requestCode,url,headers=headers)
        #f contains json as a return value
        except urllib3.exceptions.HTTPError as e:
            print(e)
            return e
    
    def mkcol(self,url):
        return self.sendRequest(url,'MKCOL')

    def delete(self,url):
        return self.sendRequest(url,'DELETE')

    def propfind(self,url,PROPFIND):
        headers=urllib3.util.make_headers(basic_auth=self.getBasicAuth())
        headers["Content-Type"]='text/xml; charset="utf-8"'
        headers['content-length']=str(len(PROPFIND))
        headers['Depth']='1'
        #add csrf
        headers["X-LABKEY-CSRF"]=self.getCSRF()
        headers["Cookie"]=self.cookie

        try:
            return self.http.request('PROPFIND',url,headers=headers,body=PROPFIND)
        #f contains json as a return value
        except urllib3.exceptions.HTTPError as e:
            print(e)
            return e

    @staticmethod 
    def HTTPStatus(response,method=None):
        if response.status==200:
            return True
        if method=='propfind':
            if response.status==207:
                return True
        
        print("Status: {}".format(response.status))
        print("Data: {}".format(response.data))
        return False

    #a good testing routine; CSRF is required for content modifying requests
    def getCSRF(self):
        url=self.GetLabkeyUrl()+'/login/whoAmI.view'
        try:
            response=self.get(url)
            encoding=chardet.detect(response.data)['encoding']
            jsonData=json.loads(response.data.decode(encoding))
        except AttributeError:
            print("Failed")
            return None
        #local cookie jar
        self.cookie=response.getheader('Set-Cookie')
        print("CSRF: {}".format(jsonData["CSRF"]))
        return jsonData["CSRF"]


    #file manipulation routiens

    #was GetFile
    def DownloadFileToCache(self,relativePath):
        debug=True
        # 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)

        #make sure we are at the begining of the file


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

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

        with open(localPath,'wb') as localBuffer:
            shutil.copyfileobj(remoteBuffer,localBuffer)


        print("Local size: {}".format(os.path.getsize(localPath)))

        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 entryExists(self,url):

        #use head
        response=self.head(url)
        return labkeyURIHandler.HTTPStatus(response)

    def remoteDirExists(self,url):
        #weaker, just checks if there is an entry at url, does not check wheter it is a dir
        return self.entryEists(url)
        #stronger, but more complicated
        return self.isRemoteDir(url)

    def mkdir(self,remoteDir):
        if self.remoteDirExists(remoteDir):
            return False
        response=self.mkcol(remoteDir)
        return labkeyURIHandler.HTTPStatus(response)

    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
        return labkeyURIHandler.HTTPStatus(self.delete(remoteDir))

    #was listDir
    def listRelativeDir(self,relativeDir):
        print("Listing for {0}".format(relativeDir))
        dirUrl=self.GetLabkeyPathFromRelativePath(relativeDir)
        print("Setting url: {}".format(dirUrl))
        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)
        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>"""
        if not labkeyURIHandler.HTTPStatus(self.propfind(remotePath,PROPFIND),method='propfind'):
            print("Bad status")
            return False
        
        try:
            tree=ET.XML(f.data)
        except ET.ParseError:
            #print(f.data.decode('utf-8'))
            #if directory is not there, a 404 response will be made by HTTP, which is not an xml file
            #we are safe to assume the directory is not there
            return False
        
        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

        PROPFIND=u"""<?xml version="1.0" encoding="utf-8"?>\n
                    <propfind xmlns="DAV:">\n
                    <prop>\n
                    <getetag/>\n
                    </prop>\n
                    </propfind>"""
        response=self.propfind(dirUrl,PROPFIND)
        if not labkeyURIHandler.HTTPStatus(response,method='propfind'):
            print("Bad status")
            return False,[]

        try:
            tree=ET.XML(response.data)
        except ET.ParseError:
            print("Fail to parse XML")
            return False,[]

        rps=tree.findall('{DAV:}response')
        dirs=[]
        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)
        response=self.get(dirUrl,binary=True)
        if not labkeyURIHandler.HTTPStatus(response):
            return io.BytesIO()
        #this will collect for all data
        remoteBuffer=io.BytesIO(response.data)
        response.release_conn()
        return remoteBuffer

    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()
        self.put(remotePath,data)

    #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
    @staticmethod
    def getJSON(response):
        encoding=chardet.detect(response.data)['encoding']
        return json.loads(response.data.decode(encoding))


    def loadDataset(self,project,dataset):
        url=self.GetLabkeyUrl()+'/'+project
        url+='/query-selectRows.api?schemaName=study&query.queryName='+dataset
        return labkeyURIHandler.getJSON(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 labkeyURIHandler.getJSON(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)).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 labkeyURIHandler.getJSON(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)).data