#!/usr/bin/env python
import sys
import block

f = open(sys.argv[1],'r')

volumeid = f.read(512)

def getUInt8(data,offset):
	return block.Readview('B',block.Block(data[offset:offset+1]))[0]

def getUInt16(data,offset):
	return block.Readview('H',block.Block(data[offset:offset+2]))[0]

def getUInt32(data,offset):
	return block.Readview('I',block.Block(data[offset:offset+4]))[0]

def syswrite(text):
	sys.stdout.write(text)
	sys.stdout.flush()


OEMIdent = volumeid[0x03:0x0A]
bytesPerSector = getUInt16(volumeid,0x0B)
sectorsPerCluster = getUInt8(volumeid,0x0D)
numReservedSectors = getUInt16(volumeid,0x0E)
numFATs = getUInt8(volumeid,0x10)
numDirEnt = getUInt16(volumeid,0x11)
totalSectors16 = getUInt16(volumeid,0x13)
totalSectors32 = getUInt32(volumeid,0x20)
sectorsPerFAT = getUInt32(volumeid,0x24)
rootDirFirstCluster = getUInt32(volumeid,0x2C)
signature = getUInt16(volumeid,0x1FE)

print """OEMIdent: %s
BytesPerSector: %i
SectorsPerCluster: %i
NumReservedSectors: 0x%X
numFATs: %i
numDirEnt: %i
totalSectors16: %i
totalSectors32: %i
sectorsPerFAT: %i
rootDirFirstCluster: 0x%.8X
signature: 0x%.4X (should be 0xAA55)
""" % (OEMIdent,bytesPerSector,sectorsPerCluster,numReservedSectors,numFATs,numDirEnt,totalSectors16,totalSectors32,sectorsPerFAT,rootDirFirstCluster,signature)

fat1addr = numReservedSectors*bytesPerSector
fat1length = sectorsPerFAT*bytesPerSector

print """------
fat1addr: \t%8i \t0x%.8X
fat1length: \t%8i \t0x%.8X
fat1:     \t0x%.8x  - \t0x%.8x
""" % (fat1addr,fat1addr,fat1length,fat1length,fat1addr,fat1addr+fat1length)

f.seek(fat1addr)
fat1 = f.read(fat1length)
fat2 = f.read(fat1length) # FIXME: there are not necessarily two

if (fat1 == fat2):
	print "FATs are the same"

fatptr = 0
usedclusters = {}

print "Finding non-zero clusters"
while (fatptr < fat1length/4): # 256 now
	nextlink = getUInt32(fat1,fatptr*4)
	if (nextlink > 0): # if this block does something
		usedclusters[fatptr] = nextlink
	fatptr += 1
	if ((fatptr % 1024) == 0):
		syswrite(".")

print "\nFound %i used clusters = %iMB data\nTraversing files..." % (len(usedclusters),(len(usedclusters)*sectorsPerCluster*bytesPerSector)/1048576)

files = []

while True:
	keys = usedclusters.keys()
	if (len(keys) > 0):
		nextptr = usedclusters.keys()[0];
	else:
		break
	newfile = []
	while True:
		thisptr = nextptr
		if not (usedclusters.has_key(nextptr)):
			print "Error: nextptr %.8X is unallocated" % nextptr
			break
		nextptr = usedclusters[nextptr]
		del usedclusters[thisptr]
		newfile.append(thisptr)
		if (nextptr >= 0x0FFFFFF8):
			break
	files.append(newfile)

print "%i files found" % len(files)

for i in files:
	print "File found: %i clusters (%i MB), start %.8X" % (len(i),len(i)*sectorsPerCluster*bytesPerSector/1048576,i[0])

firstDataSector = (numReservedSectors * bytesPerSector) + (numFATs * fat1length)
print "firstDataSector %i 0x%.8X" % (firstDataSector,firstDataSector) 

def getCluster(cluster):
	address = (((cluster - 2) * sectorsPerCluster) * bytesPerSector) + firstDataSector
	f.seek(address)
	return f.read(bytesPerSector * sectorsPerCluster)

index = 0
for file in files:
	g = open("recovered/%i" % index,'w')
	for cluster in file:
		g.write(getCluster(cluster))
	g.close()
	index += 1
	print "written %i" % (index)