Panda3D Multifiles implemented in Python

multifile.py

from panda3d.core import IDecryptStream, StringStream
from panda3d.core import Datagram, DatagramIterator

SF_compressed = 0x0008
SF_encrypted = 0x0010
SF_signature = 0x0020

class Subfile(object):

    def __init__(self, mf):
        self.mf = mf
        self.address = -1
        self.length = -1
        self.flags = 0
        self.original_length = -1
        self.timestamp = 0
        self.name = None

    def load(self, di, address):
        index = di.get_current_index()

        if address > index:
            di.skip_bytes(address - index)

        next_address = di.get_uint32()
        
        if next_address == 0:
            return 0
        
        self.address = di.get_uint32()
        self.length = di.get_uint32()
        self.flags = di.get_uint16()

        if (self.flags & (SF_compressed | SF_encrypted)) != 0:
            self.original_length = di.get_uint32()
        else:
            self.original_length = self.length

        if self.mf.minor_version < 1:
            self.timestamp = self.mf.timestamp
        else:
            self.timestamp = di.get_uint32()

            if self.timestamp == 0:
                self.timestamp = self.mf.timestamp

        length = di.get_uint16()
        self.name = di.extract_bytes(length)
        self.name = bytes([b ^ 0xFF for b in self.name])
        return next_address

    def is_compressed(self):
        return self.flags & SF_compressed != 0

    def is_encrypted(self):
        return self.flags & SF_encrypted != 0

    def is_signature(self):
        return self.flags & SF_signature != 0

    def __str__(self):
        return 'Subfile {0} created at {1} with length {2} original length {3} flags {4} at {5}. Compressed: {6}, encrypted: {7}'.format(
            b'SIGNATURE' if self.is_signature() else self.name,
            self.timestamp,
            self.length,
            self.original_length,
            self.flags,
            self.address,
            'YES' if self.is_compressed() else 'NO',
            'YES' if self.is_encrypted() else 'NO'
        )

    def read(self, dg):
        if self.is_compressed():
            raise Exception('Not implemented')

        di = DatagramIterator(dg)
        di.skip_bytes(self.address)

        return di.extract_bytes(self.length)

class Multifile(object):
    HEADER = b'pmf\0\n\r'

    def __init__(self):
        self.major_version = 0
        self.minor_version = 0
        self.scale_factor = 0
        self.timestamp = 0

    def load(self, f):
        dg = Datagram(f.read())
        di = DatagramIterator(dg)
        header = di.extract_bytes(6)

        if header != self.HEADER:
            raise Exception('Invalid multifile header.')

        self.major_version = di.get_int16()
        self.minor_version = di.get_int16()
        self.scale_factor = di.get_uint32()
        self.timestamp = di.get_uint32()

        next_address = di.get_current_index()
        self.subfiles = []

        count = 0

        while next_address != 0:
            count += 1
            subfile = Subfile(self)
            
            next_address = subfile.load(di, next_address)

            if next_address != 0:
                self.subfiles.append(subfile)

        self.subfile = min(self.subfiles, key=lambda subfile: subfile.length)
        self.subfile_data = self.subfile.read(dg)

    def is_password(self, password):
        ss = StringStream(self.subfile_data)
        ss2 = IDecryptStream(ss, False, password)
        return ss2.read(6) == b'crypty'

    def __str__(self):
        return 'Panda3D Multifile version {0}.{1} with scale factor {2} and timestamp {3}'.format(
            self.major_version, self.minor_version,
            self.scale_factor,
            self.timestamp
        )