b2root.py

#!/usr/bin/env python3
 
 
 
import os
import sys
import tempfile
import shutil
from basf2 import B2ERROR
 
 
byteorder = 'big'
 
 
class RawRootFile:
    """
    Interface to binary root file content
    """
 
    def __init__(self, filename):
        """
        Open a root file and read its header
        """
 
        
        self.rootfilerootfile = open(filename, 'rb')
        root = self.rootfilerootfile.read(4)
        if root != b'root':
            raise Exception("%s is not a root file" % filename)
        
        self.versionversion = int.from_bytes(self.rootfilerootfile.read(4), byteorder)
        
        self.largelarge = (self.versionversion >= 1000000)
        
        self.wordlenwordlen = 8 if self.largelarge else 4
        
        self.beginbegin = int.from_bytes(self.rootfilerootfile.read(4), byteorder)
        
        self.endend = int.from_bytes(self.rootfilerootfile.read(self.wordlenwordlen), byteorder)
        
        self.seekfreeseekfree = int.from_bytes(self.rootfilerootfile.read(self.wordlenwordlen), byteorder)
        
        self.nbytesfreenbytesfree = int.from_bytes(self.rootfilerootfile.read(4), byteorder)
        
        self.nfreenfree = int.from_bytes(self.rootfilerootfile.read(4), byteorder)
        
        self.nbytesnamenbytesname = int.from_bytes(self.rootfilerootfile.read(4), byteorder)
        
        self.unitsunits = int.from_bytes(self.rootfilerootfile.read(1), byteorder)
        
        self.compresscompress = int.from_bytes(self.rootfilerootfile.read(4), byteorder)
        
        self.seekinfoseekinfo = int.from_bytes(self.rootfilerootfile.read(self.wordlenwordlen), byteorder)
        
        self.nbytesinfonbytesinfo = int.from_bytes(self.rootfilerootfile.read(self.wordlenwordlen), byteorder)
        
        self.uuiduuid = self.rootfilerootfile.read(18)
        self.rootfilerootfile.seek(0)
        
        self.headerheader = self.rootfilerootfile.read(self.beginbegin)
        self.rootfilerootfile.seek(self.beginbegin)
        
        self.seekkeysseekkeys = 0
 
    def __del__(self):
        """
        Close file
        """
 
        if hasattr(self, 'rootfile'):
            self.rootfilerootfile.close()
 
    def normalize(self, end=None, seekfree=None, nbytesfree=None, nbytesname=None, seekinfo=None):
        """
        Set UUID to zero and adjust pointers that are given as arguments
        """
 
        if end:
            self.endend = end
        if seekfree:
            self.seekfreeseekfree = seekfree
        if nbytesfree:
            self.nbytesfreenbytesfree = nbytesfree
        if nbytesname:
            self.nbytesnamenbytesname = nbytesname
        if seekinfo:
            self.seekinfoseekinfo = seekinfo
        self.uuiduuid = b'\x00' * 18
 
        self.headerheader = b'root' \
            + self.versionversion.to_bytes(4, byteorder) \
            + self.beginbegin.to_bytes(4, byteorder) \
            + self.endend.to_bytes(self.wordlenwordlen, byteorder) \
            + self.seekfreeseekfree.to_bytes(self.wordlenwordlen, byteorder) \
            + self.nbytesfreenbytesfree.to_bytes(4, byteorder) \
            + self.nfreenfree.to_bytes(4, byteorder) \
            + self.nbytesnamenbytesname.to_bytes(4, byteorder) \
            + self.unitsunits.to_bytes(1, byteorder) \
            + self.compresscompress.to_bytes(4, byteorder) \
            + self.seekinfoseekinfo.to_bytes(self.wordlenwordlen, byteorder) \
            + self.nbytesinfonbytesinfo.to_bytes(self.wordlenwordlen, byteorder) \
            + self.uuiduuid
        nzero = self.beginbegin - len(self.headerheader)
        self.headerheader += b'\x00' * nzero
 
    def __iter__(self):
        """
        iterator interface
        """
 
        return self
 
    class Key:
        """
        Root file key
        """
 
        def __init__(self, keydata, large=False):
            """
            Obtain key header and data payload from given data and extract header information
            """
 
            
            self.largelarge = large
            
            self.wordlenwordlen = 8 if self.largelarge else 4
            pos = 0
            
            self.nbytes, pos = self.get_intget_int(keydata, pos)
            
            self.version, pos = self.get_intget_int(keydata, pos, 2)
            
            self.objlen, pos = self.get_intget_int(keydata, pos)
            
            self.datimedatime, pos = keydata[pos:pos+4], pos+4
            
            self.keylen, pos = self.get_intget_int(keydata, pos, 2)
            
            self.cycle, pos = self.get_intget_int(keydata, pos, 2)
            
            self.seekkeyseekkey, pos = self.get_intget_int(keydata, pos, self.wordlenwordlen)
            
            self.seekpdir, pos = self.get_intget_int(keydata, pos, self.wordlenwordlen)
            
            self.classname, pos = self.get_stringget_string(keydata, pos)
            
            self.name, pos = self.get_stringget_string(keydata, pos)
            
            self.title, pos = self.get_stringget_string(keydata, pos)
            
            self.headerheader = keydata[:pos]
            
            self.datadata = keydata[pos:]
            
            self.shownameshowname = self.classname
 
        def normalize(self, pos=None, offset=None):
            """
            Set the key datime to zero adjust the pointer to itself if given as argument
            """
 
            
            self.datimedatime = b'\x00' * 4
            if self.seekkeyseekkey > 0:
                if pos:
                    
                    self.seekkeyseekkey = pos
                elif offset:
                    self.seekkeyseekkey += offset
            self.recreate_headerrecreate_header()
 
        def recreate_header(self):
            """
            Build the binary header information from the data members
            """
 
            self.headerheader = self.nbytes.to_bytes(4, byteorder) \
                + self.version.to_bytes(2, byteorder) \
                + self.objlen.to_bytes(4, byteorder) \
                + self.datimedatime \
                + self.keylen.to_bytes(2, byteorder) \
                + self.cycle.to_bytes(2, byteorder) \
                + self.seekkeyseekkey.to_bytes(self.wordlenwordlen, byteorder) \
                + self.seekpdir.to_bytes(self.wordlenwordlen, byteorder) \
                + len(self.classname).to_bytes(1, byteorder) + self.classname \
                + len(self.name).to_bytes(1, byteorder) + self.name \
                + len(self.title).to_bytes(1, byteorder) + self.title
 
        def get_int(self, data, pos, wordlen=4):
            """
            Helper function to read an int from binary data
            """
 
            return (int.from_bytes(data[pos:pos+wordlen], byteorder), pos+wordlen)
 
        def get_string(self, data, pos):
            """
            Helper function to read a string from binary data
            """
 
            strlen = data[pos]
            return (data[pos+1:pos+1+strlen], pos+1+strlen)
 
    def __next__(self):
        """
        Iterate to next key
        """
 
        pos = self.rootfilerootfile.tell()
        large = pos > 0x80000000
        nbytes = int.from_bytes(self.rootfilerootfile.read(4), byteorder)
        if (nbytes == 0):
            raise StopIteration()
        self.rootfilerootfile.seek(-4, 1)
        result = self.KeyKey(self.rootfilerootfile.read(nbytes), large)
        if pos == self.seekfreeseekfree:
            result.showname = b'FreeSegments'
        elif pos == self.seekinfoseekinfo:
            result.showname = b'StreamerInfo'
        elif pos == self.seekkeysseekkeys:
            result.showname = b'KeysList'
 
        if result.showname == b'TFile':
            result.filename, pos = result.get_string(result.data, 0)
            result.filetitle, pos = result.get_string(result.data, pos)
            result.version = int.from_bytes(result.data[pos:pos+2], byteorder)
            result.nbyteskeys = int.from_bytes(result.data[pos+10:pos+14], byteorder)
            result.nbytesname = int.from_bytes(result.data[pos+14:pos+18], byteorder)
            if result.version > 1000:
                result.seekdir = int.from_bytes(result.data[pos+18:pos+26], byteorder)
                result.seekparent = int.from_bytes(result.data[pos+26:pos+34], byteorder)
                self.seekkeysseekkeys = int.from_bytes(result.data[pos+34:pos+42], byteorder)
            else:
                result.seekdir = int.from_bytes(result.data[pos+18:pos+22], byteorder)
                result.seekparent = int.from_bytes(result.data[pos+22:pos+26], byteorder)
                self.seekkeysseekkeys = int.from_bytes(result.data[pos+26:pos+30], byteorder)
        return result
 
 
def normalize_file(filename, output=None, in_place=False, name=None, root_version=None):
    """
    Reset the non-reproducible root file metadata: UUID and datimes.
    It can also reset the initial file name stored in the file itself, but
    (WARNING!) this may corrupt the root file.
    """
 
    # open input file
    rootfile = RawRootFile(filename)
 
    # adjust root version number
    if root_version:
        rootfile.version = root_version
 
    # create output file
    if output:
        newrootfile = open(output, 'wb')
    elif in_place:
        newrootfile = tempfile.TemporaryFile()
    else:
        basename, ext = os.path.splitext(filename)
        newrootfile = open(basename + '_normalized' + ext, 'wb')
 
    # write output file header
    newrootfile.write(rootfile.header)
 
    # file name in the metadata
    if name:
        newname = name.encode()
    else:
        newname = None
 
    # bookkeeping of offsets, positions, and keys
    offset = 0
    seekfree = rootfile.seekfree
    nbytesfree = rootfile.nbytesfree
    nbytesname = rootfile.nbytesname
    seekinfo = rootfile.seekinfo
    keylist = []
    keyskey = None
    infokey = None
    swap = False
 
    for key in rootfile:
        # reset datime and adjust position of key
        key.normalize(pos=newrootfile.tell())
 
        # Special treatment of key containing the TFile information
        if key.showname == b'TFile':
 
            # if a new name is given change the name and determine the offset caused by the change of name length
            namelen = len(key.name)
            if newname:
                key.name = key.filename = newname
            offset = len(key.name) - namelen
 
            # apply the offset to total (2x), object, and key length
            key.nbytes += 2*offset
            key.objlen += offset
            key.keylen += offset
 
            # recreate the header from updated data members
            key.recreate_header()
 
            # recreate the key data with new name and lengths
            buffer = len(key.name).to_bytes(1, byteorder) + key.name
            buffer += len(key.title).to_bytes(1, byteorder) + key.title
            buffer += key.version.to_bytes(2, byteorder)
            buffer += (0).to_bytes(8, byteorder)  # reset datimeC and datimeM
            buffer += (key.nbyteskeys + offset).to_bytes(4, byteorder)
            buffer += (key.nbytesname + 2*offset).to_bytes(4, byteorder)
            wordlen = 8 if key.version > 1000 else 4
            buffer += key.seekdir.to_bytes(wordlen, byteorder)
            buffer += key.seekparent.to_bytes(wordlen, byteorder)
            seekkeyspos = newrootfile.tell() + len(key.header) + len(buffer)
            buffer += (rootfile.seekkeys + 2*offset).to_bytes(wordlen, byteorder)
            buffer += (0).to_bytes(18, byteorder)  # reset UUID
            if key.version <= 1000:
                buffer += (0).to_bytes(12, byteorder)
            key.data = buffer
 
        else:
            # check whether we break pointers in TTrees
            if key.classname == b'TTree' and offset != 0:
                B2ERROR('Changing the name of root files containing a tree is not supported.')
                if not in_place:
                    os.remove(newrootfile.name)
                sys.exit(1)
 
            # update key data in KeysList: number of keys and key headers
            # and remember key position
            if key.showname == b'KeysList':
                seekkeys = newrootfile.tell()
                buffer = len(keylist).to_bytes(4, byteorder)
                for filekey in keylist:
                    buffer += filekey.header
                key.data = buffer
                keyskey = key
                swap = (infokey is None)
 
            # update free segments pointer and remember key position
            if key.showname == b'FreeSegments':
                seekfree = newrootfile.tell()
                pointer = int.from_bytes(key.data[2:6], byteorder) + 4*offset
                key.data = key.data[:2] + pointer.to_bytes(4, byteorder) + key.data[6:]
 
            # update name in KeysList and FreeSegments
            if key.showname in [b'KeysList', b'FreeSegments'] and newname:
                key.name = newname
                key.nbytes += offset
                key.keylen += offset
                key.recreate_header()
 
            # keep track of all keys for the KeysList
            elif key.showname not in [b'StreamerInfo', b'']:
                keylist.append(key)
 
            # remember streamer info key and position
            if key.showname == b'StreamerInfo':
                seekinfo = newrootfile.tell()
                infokey = key
 
        # write the updated key, making sure the KeysList come after the StreamerInfo
        if swap and key.showname == b'KeysList':
            pass
        elif swap:
            seekinfo = newrootfile.tell()
            infokey.normalize(pos=seekinfo)
            newrootfile.write(infokey.header)
            newrootfile.write(infokey.data)
            seekkeys = newrootfile.tell()
            keyskey.normalize(pos=seekkeys)
            newrootfile.write(keyskey.header)
            newrootfile.write(keyskey.data)
            swap = False
        else:
            newrootfile.write(key.header)
            newrootfile.write(key.data)
 
    # write the new file header
    rootfile.normalize(end=newrootfile.tell(), seekfree=seekfree, nbytesfree=nbytesfree+offset,
                       nbytesname=nbytesname+2*offset, seekinfo=seekinfo)
    newrootfile.seek(0)
    newrootfile.write(rootfile.header)
 
    # update pointer to keyslist
    newrootfile.seek(seekkeyspos)
    newrootfile.write(seekkeys.to_bytes(wordlen, byteorder))
 
    # replace in the input file if the in-place option is used
    if in_place:
        del rootfile
        newrootfile.seek(0)
        shutil.copyfileobj(newrootfile, open(filename, 'wb'))
 
    newrootfile.close()