def _read_extended(self, offset):
next_byte = int_from_byte(self._buffer[offset])
type_num = next_byte + 7
if type_num < 7:
raise InvalidDatabaseError(
'Something went horribly wrong in the decoder. An '
'extended type resolved to a type number < 8 '
'({type})'.format(type=type_num))
return type_num, offset + 1
def _resolve_data_pointer(self, pointer):
resolved = pointer - self._metadata.node_count + \
self._metadata.search_tree_size
if resolved >= self._buffer_size:
raise InvalidDatabaseError(
"The AW DB file's search tree is corrupt")
(data, _) = self._decoder.decode(resolved)
return data
def decode(self, offset):
new_offset = offset + 1
ctrl_byte = int_from_byte(self._buffer[offset])
type_num = ctrl_byte >> 5
if not type_num:
(type_num, new_offset) = self._read_extended(new_offset)
try:
decoder = self._type_decoder[type_num]
except KeyError:
raise InvalidDatabaseError('Unexpected type number ({type}) '
'encountered'.format(type=type_num))
(size, new_offset) = self._size_from_ctrl_byte(ctrl_byte, new_offset,
type_num)
return decoder(self, size, new_offset)
def _find_address_in_tree(self, packed):
bit_count = len(packed) * 8
node = self._start_node(bit_count)
node_count = self._metadata.node_count
i = 0
while i < bit_count and node < node_count:
bit = 1 & (packed[i >> 3] >> 7 - (i % 8))
node = self._read_node(node, bit)
i = i + 1
if node == node_count:
return 0, i
if node > node_count:
return node, i
raise InvalidDatabaseError('Invalid node in search tree')
def __init__(self, database, mode=MODE_AUTO):
if (mode == MODE_AUTO and mmap) or mode == MODE_MMAP:
with open(database, 'rb') as db_file:
self._buffer = mmap.mmap(db_file.fileno(),
0,
access=mmap.ACCESS_READ)
self._buffer_size = self._buffer.size()
filename = database
elif mode in (MODE_AUTO, MODE_FILE):
self._buffer = FileBuffer(database)
self._buffer_size = self._buffer.size()
filename = database
elif mode == MODE_MEMORY:
with open(database, 'rb') as db_file:
self._buffer = db_file.read()
self._buffer_size = len(self._buffer)
filename = database
elif mode == MODE_FD:
self._buffer = database.read()
self._buffer_size = len(self._buffer)
filename = database.name
else:
raise ValueError(
'Unsupported open mode ({0}). Only MODE_AUTO, MODE_FILE, '
'MODE_MEMORY and MODE_FD are supported by the pure Python '
'Reader'.format(mode))
metadata_start = self._buffer.rfind(
self._METADATA_START_MARKER, max(0,
self._buffer_size - 128 * 1024))
if metadata_start == -1:
self.close()
raise InvalidDatabaseError('Error opening database file ({0}). '
'Is this a valid AW DB file?'
''.format(filename))
metadata_start += len(self._METADATA_START_MARKER)
metadata_decoder = Decoder(self._buffer, metadata_start)
(metadata, _) = metadata_decoder.decode(metadata_start)
self._metadata = Metadata(**metadata)
self._decoder = Decoder(
self._buffer, self._metadata.search_tree_size +
self._DATA_SECTION_SEPARATOR_SIZE)
self.closed = False
def _read_node(self, node_number, index):
base_offset = node_number * self._metadata.node_byte_size
record_size = self._metadata.record_size
if record_size == 24:
offset = base_offset + index * 3
node_bytes = b'\x00' + self._buffer[offset:offset + 3]
elif record_size == 28:
offset = base_offset + 3 * index
node_bytes = bytearray(self._buffer[offset:offset + 4])
if index:
node_bytes[0] = 0x0F & node_bytes[0]
else:
middle = (0xF0 & node_bytes.pop()) >> 4
node_bytes.insert(0, middle)
elif record_size == 32:
offset = base_offset + index * 4
node_bytes = self._buffer[offset:offset + 4]
else:
raise InvalidDatabaseError(
'Unknown record size: {0}'.format(record_size))
return struct.unpack(b'!I', node_bytes)[0]
def _verify_size(self, expected, actual):
if expected != actual:
raise InvalidDatabaseError(
'The AW DB file\'s data section contains bad data '
'(unknown data type or corrupt data)')