<s> """<STR_LIT>""" <EOL> import os <EOL> import shutil <EOL> import StringIO <EOL> import zlib <EOL> try : <EOL> import urllib3 <EOL> from urllib3 import connectionpool <EOL> except ImportError : <EOL> from requests . packages import urllib3 <EOL> from requests . packages . urllib3 import connectionpool <EOL> import logging <EOL> from grr . lib import data_store <EOL> from grr . lib import utils <EOL> from grr . lib . data_stores import common <EOL> from grr . lib . rdfvalues import data_server as rdf_data_server <EOL> from grr . server . data_server import constants <EOL> from grr . server . data_server import store <EOL> from grr . server . data_server import utils as sutils <EOL> COPY_EXCEPTIONS = [ store . BASE_MAP_SUBJECT ] <EOL> MOVE_EXCEPTIONS = [ constants . TRANSACTION_FILENAME , constants . REMOVE_FILENAME ] <EOL> COMPRESSION_LEVEL = <NUM_LIT:3> <EOL> def _RecComputeRebalanceSize ( mapping , server_id , dspath , subpath ) : <EOL> """<STR_LIT>""" <EOL> total = <NUM_LIT:0> <EOL> fulldir = utils . JoinPath ( dspath , subpath ) <EOL> for comp in os . listdir ( fulldir ) : <EOL> if comp == constants . REBALANCE_DIRECTORY : <EOL> continue <EOL> path = utils . JoinPath ( fulldir , comp ) <EOL> name , unused_extension = os . path . splitext ( comp ) <EOL> if name in COPY_EXCEPTIONS : <EOL> logging . info ( "<STR_LIT>" , comp ) <EOL> continue <EOL> if os . path . isdir ( path ) : <EOL> total += _RecComputeRebalanceSize ( mapping , server_id , dspath , <EOL> utils . JoinPath ( subpath , comp ) ) <EOL> elif os . path . isfile ( path ) : <EOL> key = common . MakeDestinationKey ( subpath , name ) <EOL> where = sutils . MapKeyToServer ( mapping , key ) <EOL> if where != server_id : <EOL> logging . info ( "<STR_LIT>" , path , server_id , where ) <EOL> total += os . path . getsize ( path ) <EOL> else : <EOL> logging . info ( "<STR_LIT>" , path ) <EOL> return total <EOL> def ComputeRebalanceSize ( mapping , server_id ) : <EOL> """<STR_LIT>""" <EOL> loc = data_store . DB . Location ( ) <EOL> if not os . path . exists ( loc ) : <EOL> return <NUM_LIT:0> <EOL> if not os . path . isdir ( loc ) : <EOL> return <NUM_LIT:0> <EOL> return _RecComputeRebalanceSize ( mapping , server_id , loc , "<STR_LIT>" ) <EOL> class FileCopyWrapper ( object ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , rebalance , directory , filename , fullpath ) : <EOL> filesize = os . path . getsize ( fullpath ) <EOL> filecopy = rdf_data_server . DataServerFileCopy ( rebalance_id = rebalance . id , <EOL> directory = directory , <EOL> filename = filename , <EOL> size = filesize ) <EOL> filecopy_str = filecopy . SerializeToString ( ) <EOL> self . header = sutils . SIZE_PACKER . pack ( len ( filecopy_str ) ) <EOL> self . header += filecopy_str <EOL> self . header = StringIO . StringIO ( self . header ) <EOL> self . fp = open ( fullpath , "<STR_LIT:rb>" ) <EOL> self . compressor = zlib . compressobj ( COMPRESSION_LEVEL ) <EOL> self . buffered = "<STR_LIT>" <EOL> self . end_of_file = False <EOL> self . end_of_stream = False <EOL> self . read_header = True <EOL> def read ( self , blocksize ) : <EOL> """<STR_LIT>""" <EOL> if self . end_of_stream : <EOL> return "<STR_LIT>" <EOL> if self . read_header : <EOL> ret = self . header . read ( blocksize ) <EOL> if ret : <EOL> return ret <EOL> self . read_header = False <EOL> blocksize -= sutils . SIZE_PACKER . size <EOL> while not self . buffered and not self . end_of_file : <EOL> raw = self . fp . read ( blocksize ) <EOL> if not raw : <EOL> self . end_of_file = True <EOL> self . buffered += self . compressor . flush ( ) <EOL> break <EOL> compressed = self . compressor . compress ( raw ) <EOL> if compressed : <EOL> self . buffered += compressed <EOL> if len ( self . buffered ) > blocksize : <EOL> ret = self . buffered [ : blocksize ] <EOL> self . buffered = self . buffered [ blocksize : ] <EOL> else : <EOL> ret = self . buffered <EOL> self . buffered = "<STR_LIT>" <EOL> if not ret : <EOL> self . end_of_stream = True <EOL> return sutils . SIZE_PACKER . pack ( len ( ret ) ) + ret <EOL> def close ( self ) : <EOL> self . fp . close ( ) <EOL> self . header . close ( ) <EOL> def _SendFileToServer ( pool , fullpath , subpath , basename , rebalance ) : <EOL> """<STR_LIT>""" <EOL> fp = FileCopyWrapper ( rebalance , subpath , basename , fullpath ) <EOL> try : <EOL> headers = { "<STR_LIT>" : <NUM_LIT:0> } <EOL> res = pool . urlopen ( "<STR_LIT:POST>" , "<STR_LIT>" , headers = headers , <EOL> body = fp ) <EOL> if res . status != constants . RESPONSE_OK : <EOL> return False <EOL> except urllib3 . exceptions . MaxRetryError : <EOL> logging . warning ( "<STR_LIT>" , fullpath ) <EOL> return False <EOL> finally : <EOL> fp . close ( ) <EOL> return True <EOL> def _GetTransactionDirectory ( database_dir , rebalance_id ) : <EOL> dirname = common . ConvertStringToFilename ( rebalance_id ) <EOL> tempdir = utils . JoinPath ( database_dir , constants . REBALANCE_DIRECTORY , <EOL> dirname ) <EOL> return tempdir <EOL> def _CreateDirectory ( database_dir , rebalance_id ) : <EOL> tempdir = _GetTransactionDirectory ( database_dir , rebalance_id ) <EOL> try : <EOL> os . makedirs ( tempdir ) <EOL> except os . error : <EOL> pass <EOL> return tempdir <EOL> def _FileWithRemoveList ( database_dir , rebalance ) : <EOL> tempdir = _CreateDirectory ( database_dir , rebalance . id ) <EOL> return utils . JoinPath ( tempdir , constants . REMOVE_FILENAME ) <EOL> def _RecCopyFiles ( rebalance , server_id , dspath , subpath , <EOL> pool_cache , removed_list ) : <EOL> """<STR_LIT>""" <EOL> fulldir = utils . JoinPath ( dspath , subpath ) <EOL> mapping = rebalance . mapping <EOL> for comp in os . listdir ( fulldir ) : <EOL> if comp == constants . REBALANCE_DIRECTORY : <EOL> continue <EOL> path = utils . JoinPath ( fulldir , comp ) <EOL> name , unused_extension = os . path . splitext ( comp ) <EOL> if name in COPY_EXCEPTIONS : <EOL> continue <EOL> if os . path . isdir ( path ) : <EOL> result = _RecCopyFiles ( rebalance , server_id , dspath , <EOL> utils . JoinPath ( subpath , comp ) , pool_cache , <EOL> removed_list ) <EOL> if not result : <EOL> return False <EOL> continue <EOL> if not os . path . isfile ( path ) : <EOL> continue <EOL> key = common . MakeDestinationKey ( subpath , name ) <EOL> where = sutils . MapKeyToServer ( mapping , key ) <EOL> if where != server_id : <EOL> server = mapping . servers [ where ] <EOL> addr = server . address <EOL> port = server . port <EOL> key = ( addr , port ) <EOL> try : <EOL> pool = pool_cache [ key ] <EOL> except KeyError : <EOL> pool = connectionpool . HTTPConnectionPool ( addr , port = port ) <EOL> pool_cache [ key ] = pool <EOL> logging . info ( "<STR_LIT>" , key , server_id , where ) <EOL> if not _SendFileToServer ( pool , path , subpath , comp , rebalance ) : <EOL> return False <EOL> removed_list . append ( path ) <EOL> else : <EOL> logging . info ( "<STR_LIT>" , path ) <EOL> return True <EOL> def CopyFiles ( rebalance , server_id ) : <EOL> """<STR_LIT>""" <EOL> loc = data_store . DB . Location ( ) <EOL> if not os . path . exists ( loc ) : <EOL> return True <EOL> if not os . path . isdir ( loc ) : <EOL> return True <EOL> pool_cache = { } <EOL> removed_list = [ ] <EOL> ok = _RecCopyFiles ( rebalance , server_id , loc , "<STR_LIT>" , pool_cache , removed_list ) <EOL> if not ok : <EOL> return False <EOL> remove_file = _FileWithRemoveList ( loc , rebalance ) <EOL> with open ( remove_file , "<STR_LIT:w>" ) as fp : <EOL> for f in removed_list : <EOL> fp . write ( f . encode ( "<STR_LIT:utf8>" ) + "<STR_LIT:\n>" ) <EOL> return True <EOL> def SaveTemporaryFile ( fp ) : <EOL> """<STR_LIT>""" <EOL> loc = data_store . DB . Location ( ) <EOL> if not os . path . exists ( loc ) : <EOL> return False <EOL> if not os . path . isdir ( loc ) : <EOL> return False <EOL> filecopy_len_str = fp . read ( sutils . SIZE_PACKER . size ) <EOL> filecopy_len = sutils . SIZE_PACKER . unpack ( filecopy_len_str ) [ <NUM_LIT:0> ] <EOL> filecopy = rdf_data_server . DataServerFileCopy ( fp . read ( filecopy_len ) ) <EOL> rebdir = _CreateDirectory ( loc , filecopy . rebalance_id ) <EOL> filedir = utils . JoinPath ( rebdir , filecopy . directory ) <EOL> try : <EOL> os . makedirs ( filedir ) <EOL> except OSError : <EOL> pass <EOL> filepath = utils . JoinPath ( filedir , filecopy . filename ) <EOL> logging . info ( "<STR_LIT>" , filepath ) <EOL> with open ( filepath , "<STR_LIT:wb>" ) as wp : <EOL> decompressor = zlib . decompressobj ( ) <EOL> while True : <EOL> block_len_str = fp . read ( sutils . SIZE_PACKER . size ) <EOL> block_len = sutils . SIZE_PACKER . unpack ( block_len_str ) [ <NUM_LIT:0> ] <EOL> if not block_len : <EOL> break <EOL> block = fp . read ( block_len ) <EOL> to_decompress = decompressor . unconsumed_tail + block <EOL> while to_decompress : <EOL> decompressed = decompressor . decompress ( to_decompress ) <EOL> if decompressed : <EOL> wp . write ( decompressed ) <EOL> to_decompress = decompressor . unconsumed_tail <EOL> else : <EOL> to_decompress = "<STR_LIT>" <EOL> remaining = decompressor . flush ( ) <EOL> if remaining : <EOL> wp . write ( remaining ) <EOL> if os . path . getsize ( filepath ) != filecopy . size : <EOL> logging . error ( "<STR_LIT>" , filepath , filecopy . size ) <EOL> return False <EOL> return True <EOL> def _RecMoveFiles ( tempdir , dspath , subpath ) : <EOL> fulltempdir = utils . JoinPath ( tempdir , subpath ) <EOL> fulldsdir = utils . JoinPath ( dspath , subpath ) <EOL> try : <EOL> os . makedirs ( fulldsdir ) <EOL> except OSError : <EOL> pass <EOL> for fname in os . listdir ( fulltempdir ) : <EOL> if fname in MOVE_EXCEPTIONS : <EOL> continue <EOL> temppath = utils . JoinPath ( fulltempdir , fname ) <EOL> if os . path . isfile ( temppath ) : <EOL> newpath = utils . JoinPath ( fulldsdir , fname ) <EOL> logging . info ( "<STR_LIT>" , temppath , newpath ) <EOL> os . rename ( temppath , newpath ) <EOL> elif os . path . isdir ( temppath ) : <EOL> _RecMoveFiles ( tempdir , dspath , utils . JoinPath ( subpath , fname ) ) <EOL> def MoveFiles ( rebalance , is_master ) : <EOL> """<STR_LIT>""" <EOL> loc = data_store . DB . Location ( ) <EOL> if not os . path . exists ( loc ) : <EOL> return False <EOL> if not os . path . isdir ( loc ) : <EOL> return False <EOL> tempdir = _CreateDirectory ( loc , rebalance . id ) <EOL> remove_file = _FileWithRemoveList ( loc , rebalance ) <EOL> to_remove = [ ] <EOL> if os . path . exists ( remove_file ) : <EOL> to_remove = [ line . decode ( "<STR_LIT:utf8>" ) . rstrip ( "<STR_LIT:\n>" ) <EOL> for line in open ( remove_file , "<STR_LIT:r>" ) ] <EOL> for fname in to_remove : <EOL> if not fname . startswith ( loc ) : <EOL> logging . warning ( "<STR_LIT>" , fname ) <EOL> continue <EOL> if not os . path . exists ( fname ) : <EOL> logging . warning ( "<STR_LIT>" , fname ) <EOL> continue <EOL> if not os . path . isfile ( fname ) : <EOL> logging . warning ( "<STR_LIT>" , fname ) <EOL> continue <EOL> os . unlink ( fname ) <EOL> logging . info ( "<STR_LIT>" , fname ) <EOL> try : <EOL> os . unlink ( remove_file ) <EOL> except OSError : <EOL> pass <EOL> try : <EOL> _RecMoveFiles ( tempdir , loc , "<STR_LIT>" ) <EOL> except OSError : <EOL> return False <EOL> if not is_master : <EOL> if tempdir . startswith ( loc ) : <EOL> shutil . rmtree ( tempdir ) <EOL> return True <EOL> def SaveCommitInformation ( rebalance ) : <EOL> """<STR_LIT>""" <EOL> loc = data_store . DB . Location ( ) <EOL> if not os . path . exists ( loc ) : <EOL> return False <EOL> if not os . path . isdir ( loc ) : <EOL> return False <EOL> tempdir = _CreateDirectory ( loc , rebalance . id ) <EOL> tempfile = utils . JoinPath ( tempdir , constants . TRANSACTION_FILENAME ) <EOL> with open ( tempfile , "<STR_LIT:wb>" ) as fp : <EOL> fp . write ( rebalance . SerializeToString ( ) ) <EOL> return True <EOL> def DeleteCommitInformation ( rebalance ) : <EOL> """<STR_LIT>""" <EOL> loc = data_store . DB . Location ( ) <EOL> if not os . path . exists ( loc ) : <EOL> return False <EOL> if not os . path . isdir ( loc ) : <EOL> return False <EOL> tempdir = _GetTransactionDirectory ( loc , rebalance . id ) <EOL> tempfile = utils . JoinPath ( tempdir , constants . TRANSACTION_FILENAME ) <EOL> try : <EOL> os . unlink ( tempfile ) <EOL> except OSError : <EOL> pass <EOL> return True <EOL> def GetCommitInformation ( transid ) : <EOL> """<STR_LIT>""" <EOL> loc = data_store . DB . Location ( ) <EOL> if not os . path . exists ( loc ) : <EOL> return False <EOL> if not os . path . isdir ( loc ) : <EOL> return False <EOL> tempdir = _GetTransactionDirectory ( loc , transid ) <EOL> tempfile = utils . JoinPath ( tempdir , constants . TRANSACTION_FILENAME ) <EOL> if not os . path . exists ( tempfile ) : <EOL> return None <EOL> if not os . path . isfile ( tempfile ) : <EOL> return None <EOL> with open ( tempfile , "<STR_LIT:rb>" ) as fp : <EOL> return rdf_data_server . DataServerRebalance ( fp . read ( ) ) <EOL> def RemoveDirectory ( rebalance ) : <EOL> """<STR_LIT>""" <EOL> loc = data_store . DB . Location ( ) <EOL> if not os . path . exists ( loc ) : <EOL> return False <EOL> if not os . path . isdir ( loc ) : <EOL> return False <EOL> tempdir = _GetTransactionDirectory ( loc , rebalance . id ) <EOL> try : <EOL> if tempdir . startswith ( loc ) : <EOL> shutil . rmtree ( tempdir ) <EOL> except OSError : <EOL> pass </s>
