from module import Module
import tomb
import os
import winver
from runcmd import runProcess
from mount import mount,umount

def getInstance():
    return RegistryModule()

class RegistryModule(Module):

    def __init__(self):
        self.name = "winreg"
        self.description = "Extracts Windows Registry files"
        self.requiredVars = ["winvol"]
        self.vars = {}

    def run(self):
        path = tomb.getPath() + self.name + "/"
        os.mkdir(path)
        for vol in self.vars['winvol'].value:
            mntpoint = "/mnt/"
            mntid = mount("/dev/" + vol)
            mntpoint += mntid
            files = []
            windir = winver.getWindowsDirectory(mntpoint)
            if windir == None:
                raise Exception("No Windows installation present")
            version = winver.getWindowsVersion(mntpoint)
            profiles = winver.getUserProfiles(mntpoint)
            if version <= winver._WIN_ME:
                #9x
                files = [windir + "USER.DAT",windir + "SYSTEM.DAT"]
                if(len(profiles) > 0):
                    for profile in profiles:
                        if(os.path.isfile(profile + "USER.DAT")):
                            files += [profile + "USER.DAT"]
                if version == winver._WIN_ME:
                    #ME
                    files += [windir + "CLASSES.DAT"]

            elif version > winver._WIN_ME:
                    #NT
                    files += [windir + "/System32/config/SAM" ]
                    files += [windir + "/System32/config/SECURITY"]
                    files += [windir + "/System32/config/SOFTWARE"]
                    files += [windir + "/System32/config/SYSTEM"]
                    files += [windir + "/System32/config/DEFAULT"]
                    for profile in profiles:
                        files += [profile + "/NTUSER.DAT"]
                        if version > winver._WIN_NT and version < winver._WIN_VISTA:
                            #2k XP 2k3
                            files += [profile + "/Local Settings/Application Data/Microsoft/Windows/UsrClass.dat"]
                        else:
                            #Vista+
                            files += [profile + "/AppData/Local/Microsoft/Windows/UsrClass.dat"]

            print(files)
            runProcess(["tar","-czvf",path + "winreg_" + vol + ".tar.gz"] + files)
            umount(mntid)