hacks/dump-imessages/iphone-dataprotection/python_scripts/hfs/btree.py

269 lines
9.8 KiB
Python

from structs import *
"""
Probably buggy
HAX, only works on case SENSITIVE
"""
class BTree(object):
def __init__(self, file, keyStruct, dataStruct):
self.file = file
self.keyStruct = keyStruct
self.dataStruct = dataStruct
block0 = self.file.readBlock(0)
btnode = BTNodeDescriptor.parse(block0)
assert btnode.kind == kBTHeaderNode
self.header = BTHeaderRec.parse(block0[BTNodeDescriptor.sizeof():])
assert self.header.keyCompareType == 0 or self.header.keyCompareType == 0 or kHFSBinaryCompare
#TODO: do more testing when nodeSize != blockSize
self.nodeSize = self.header.nodeSize
self.nodesInBlock = file.blockSize / self.header.nodeSize
self.blocksForNode = self.header.nodeSize / file.blockSize
#print file.blockSize , self.header.nodeSize
self.lastRecordNumber = 0
type, (hdr, maprec) = self.readBtreeNode(0)
assert len(maprec) == self.nodeSize - 256
if self.header.totalNodes / 8 > len(maprec):
pass #TODO: handle map records
self.maprec = maprec
def isNodeInUse(self, nodeNumber):
thisByte = ord(self.maprec[nodeNumber / 8])
return (thisByte & (1 << (7 - (nodeNumber % 8)))) != 0
def readEmptySpace(self):
res = ""
z = 0
for i in xrange(self.header.totalNodes):
if not self.isNodeInUse(i):
z += 1
res += self.readNode(i)
assert z == self.header.freeNodes
return res
#convert construct structure to tuple
def getComparableKey(self, k):
raise Exception("implement in subclass")
def compareKeys(self, k1, k2):
k2 = self.getComparableKey(k2)
if k1 == k2:
return 0
return -1 if k1 < k2 else 1
def printLeaf(self, key, data):
print key, data
def readNode(self, nodeNumber):
node = ""
for i in xrange(self.blocksForNode):
node += self.file.readBlock(nodeNumber * self.blocksForNode + i)
return node
def readBtreeNode(self, nodeNumber):
self.lastnodeNumber = nodeNumber
node = self.readNode(nodeNumber)
self.lastbtnode = btnode = BTNodeDescriptor.parse(node)
if btnode.kind == kBTHeaderNode:
assert btnode.numRecords == 3
end = self.nodeSize - 8 #2*4
offsets = Array(btnode.numRecords+1, UBInt16("off")).parse(node[end:])
assert offsets[-4] == end
hdr = BTHeaderRec.parse(node[BTNodeDescriptor.sizeof():])
maprec = node[offsets[-3]:end]
return kBTHeaderNode, [hdr, maprec]
elif btnode.kind == kBTIndexNode:
recs = []
offsets = Array(btnode.numRecords, UBInt16("off")).parse(node[-2*btnode.numRecords:])
for i in xrange(btnode.numRecords):
off = offsets[btnode.numRecords-i-1]
k = self.keyStruct.parse(node[off:])
off += 2 + k.keyLength
k.childNode = UBInt32("nodeNumber").parse(node[off:off+4])
recs.append(k)
return kBTIndexNode, recs
elif btnode.kind == kBTLeafNode:
recs = []
offsets = Array(btnode.numRecords, UBInt16("off")).parse(node[-2*btnode.numRecords:])
for i in xrange(btnode.numRecords):
off = offsets[btnode.numRecords-i-1]
k = self.keyStruct.parse(node[off:])
off += 2 + k.keyLength
d = self.dataStruct.parse(node[off:])
recs.append((k,d))
return kBTLeafNode, recs
else:
raise Exception("Invalid node type " + str(btnode))
def search(self, searchKey, node=None):
if node == None:
node = self.header.rootNode
type, stuff = self.readBtreeNode(node)
if len(stuff) == 0:
return None, None
if type == kBTIndexNode:
for i in xrange(len(stuff)):
if self.compareKeys(searchKey, stuff[i]) < 0:
if i > 0:
i = i - 1
return self.search(searchKey, stuff[i].childNode)
return self.search(searchKey, stuff[len(stuff)-1].childNode)
elif type == kBTLeafNode:
self.lastRecordNumber = 0
for k,v in stuff:
res = self.compareKeys(searchKey, k)
if res == 0:
return k, v
if res < 0:
return None, None
self.lastRecordNumber += 1
return None, None
def traverse(self, node=None, count=0, callback=None):
if node == None:
node = self.header.rootNode
type, stuff = self.readBtreeNode(node)
if type == kBTIndexNode:
for i in xrange(len(stuff)):
count += self.traverse(stuff[i].childNode, callback=callback)
elif type == kBTLeafNode:
for k,v in stuff:
if callback:
callback(k,v)
else:
self.printLeaf(k, v)
count += 1
return count
def traverseLeafNodes(self, callback=None):
nodeNumber = self.header.firstLeafNode
count = 0
while nodeNumber != 0:
_, stuff = self.readBtreeNode(nodeNumber)
count += len(stuff)
for k,v in stuff:
if callback:
callback(k,v)
else:
self.printLeaf(k, v)
nodeNumber = self.lastbtnode.fLink
return count
#XXX
def searchMultiple(self, searchKey, filterKeyFunction=lambda x:False):
self.search(searchKey)
nodeNumber = self.lastnodeNumber
recordNumber = self.lastRecordNumber
kv = []
while nodeNumber != 0:
_, stuff = self.readBtreeNode(nodeNumber)
for k,v in stuff[recordNumber:]:
if filterKeyFunction(k):
kv.append((k,v))
else:
return kv
nodeNumber = self.lastbtnode.fLink
recordNumber = 0
return kv
def getLBAsHax(self):
nodes = [self.lastnodeNumber]
n = self.lastbtnode
for i in xrange(2):
nodes.append(self.lastbtnode.bLink)
self.readBtreeNode(self.lastbtnode.bLink)
self.lastbtnode = n
for i in xrange(2):
nodes.append(self.lastbtnode.fLink)
self.readBtreeNode(self.lastbtnode.fLink)
res = []
for n in nodes:
res.append(self.file.getLBAforBlock(n * self.blocksForNode))
return res
class CatalogTree(BTree):
def __init__(self, file, volume):
super(CatalogTree,self).__init__(file, HFSPlusCatalogKey, HFSPlusCatalogData)
self.volume = volume
def printLeaf(self, k, d):
if d.recordType == kHFSPlusFolderRecord or d.recordType == kHFSPlusFileRecord:
print getString(k)
def getComparableKey(self, k2):
#XXX http://dubeiko.com/development/FileSystems/HFSPLUS/tn1150.html#StringComparisonAlgorithm
return (k2.parentID, getString(k2))
def searchByCNID(self, cnid):
threadk, threadd = self.search((cnid, ""))
return self.search((threadd.data.parentID, getString(threadd.data))) if threadd else (None, None)
def getFolderContents(self, cnid):
return self.searchMultiple((cnid, ""), lambda k:k.parentID == cnid)
def getRecordFromPath(self, path):
if not path.startswith("/"):
return None, None
if path == "/":
return self.searchByCNID(kHFSRootFolderID)
parentId=kHFSRootFolderID
i = 1
k, v = None, None
for p in path.split("/")[1:]:
if p == "":
break
k,v = self.search((parentId, p))
if (k,v) == (None, None):
return None, None
if v.recordType == kHFSPlusFolderRecord:
parentId = v.data.folderID
elif v.recordType == kHFSPlusFileRecord and is_symlink(v.data):
linkdata = self.volume.readFileByRecord(v)
print "symlink %s => %s" % (p, linkdata)
if not linkdata:
return None, None
t = path.split("/")
t[i] = linkdata
newpath = "/".join(t)
return self.getRecordFromPath(newpath)
else:
break
i += 1
return k,v
class ExtentsOverflowTree(BTree):
def __init__(self, file):
super(ExtentsOverflowTree,self).__init__(file, HFSPlusExtentKey, HFSPlusExtentRecord)
def getComparableKey(self, k2):
return (k2.fileID, k2.forkType, k2.startBlock)
def searchExtents(self, fileID, forkType, startBlock):
return self.search((fileID, forkType, startBlock))
class AttributesTree(BTree):
def __init__(self, file):
super(AttributesTree,self).__init__(file, HFSPlusAttrKey, HFSPlusAttrData)
def printLeaf(self, k, d):
print k.fileID, getString(k), d.data.encode("hex")
def getComparableKey(self, k2):
return (k2.fileID, getString(k2))
def searchXattr(self, fileID, name):
k,v = self.search((fileID, name))
return v.data if v else None
def getAllXattrs(self, fileID):
res = {}
for k,v in self.searchMultiple((fileID, ""), lambda k:k.fileID == fileID):
res[getString(k)] = v.data
return res