import os
import sys
from hfs.emf import *
from hfs.journal import *

def do_emf_carving(volume, carveokdir, carvenokdir):
    deletedFiles, filekeys = carveEMFVolumeJournal(volume)
    print "Journal carving done, trying to extract deleted files"
    n = 0
    for name, vv in deletedFiles:
        for filekey in filekeys.get(vv.data.fileID, []):
            ff = EMFFile(volume,vv.data.dataFork, vv.data.fileID, filekey, deleted=True)
            data = ff.readAllBuffer()
            if isDecryptedCorrectly(data):
                write_file(carveokdir + "%s_%s" % (filekey.encode("hex")[:8],name.replace("/","_")),data)
                n += 1
            else:
                write_file(carvenokdir + "%s_%s" % (filekey.encode("hex")[:8],name.replace("/","_")),data)
        if not filekeys.has_key(vv.data.fileID):
            print "Missing file key for", name
        else:
            del filekeys[vv.data.fileID]
    
    print "Done, extracted %d files" % n

    if True:
        fks = set(reduce(lambda x,y: x+y, filekeys.values()))
        carveEMFemptySpace(volume, fks, carvenokdir)

if __name__ == "__main__":
    if len(sys.argv) < 2:
        print "Usage: emf_undelete.py disk_image.bin"
        sys.exit(0)
    filename = sys.argv[1]
    volume = EMFVolume(filename)
    dirname = os.path.dirname(filename)
    if dirname == "":
        dirname = "."
    outdir = dirname + "/" + volume.volumeID().encode("hex") + "_" + os.path.basename(filename)
    carveokdir = outdir + "/undelete/"
    carvenokdir = outdir + "/junk/"
    try:
        os.makedirs(carveokdir)
        os.makedirs(carvenokdir)
    except:
        pass

    do_emf_carving(volume, carveokdir, carvenokdir)