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(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"

    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):
        debug=False
        if debug:
            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
        if debug:
	         remote.seek(0,2)
	         sz=remote.tell()
	         print("Remote size: {0}").format(sz)

        remote.seek(0)
        shutil.copyfileobj(remote,local)
        if debug:
	         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 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):
        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)
        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):

        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)
        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")
        #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)
        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):

        print("PUT: {0}").format(url)
        r=MethodRequest(url.encode('utf-8'),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 getCSRF(self):
        url=self.GetLabkeyUrl()+'/login/whoAmI.view'
        jsonData=json.load(self.get(url))
        return jsonData["CSRF"]


    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 isDir(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):
        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 copyFileToRemote(self,localPath, remotePath):
        #get upstream directories sorted out

        labkeyDir=os.path.dirname(remotePath)
        if not self.remoteDirExists(labkeyDir):
            if not self.mkdirs(labkeyDir):
                errorCode="UploadFile: Could not create directory {}"
                print(errorCode.format(labkeyDir))
                return False

        #make an URL request
        if (self.isDir(remotePath)):
            remotePath=remotePath+'/'+os.path.basename(localPath)

        with open(localPath, 'r') as f:
            data=f.read()
        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
                print("self.readDir(f) not implemented")
        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 filterDataset(self,project,dataset,variable,value,oper="eq"):
        debug=True
        url=self.GetLabkeyUrl()+'/'+project
        url+='/query-selectRows.api?schemaName=study&query.queryName='+dataset
        url+="&query."+variable+"~"+oper+"="+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))