#!/usr/bin/env python
# -*- Mode: Python; tab-width: 4 -*-
#
# inctrl2pe.py - Inctrl5 csv report to pe inf format converter
# Copyright (C) 2004 Sherpya <sherpya@hotmail.com>
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by the
# Free Software Foundation; either version 2, or (at your option) any later
# version.
#
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTIBILITY
# or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License
# for more details.
# ======================================================================
__version__ = '0.6'
__pychecker__ = 'maxreturns=11'

### History
### Version 0.1:
### - First public release, it parses registry entries with a "smart engine"
### Version 0.2:
### - Basic support for files entries - it has some bug ;( - Unreleased
### Version 0.3:
### - Added support for "real" csv (using ',' instead of ';')
### - Blacklist works for files and directories
### - Fixed (I hope) file handling
### Version 0.4:
### - Tweaked some blacklist keys
### - Added a replacelist like in reg2pe
### - Win32: Suggest basename.inf file on save
### - Win32: Added auto-open conf value (with ShellExecute)
### - Added NOPATH to strip path from InprocServer32 keys
### Version 0.5:
### - 0x1,"Something", -> 0x0,"Something"
### - Added ONLYADDED option to log only new added values
### - Rewritten splitting routine, it was causing problems with @some,-1234 values
### - Added CLEANUP option to remove uneeded keys
### - HELPDIR and NOPATH fixes
### Version 0.6
### - Make pychecker happy
### - Code cleanup
### - Fixed a bug when decoding MULTI_SZ

from sys import argv, stderr, platform, exit as sys_exit
from string import letters
import re

### Constants

tkeys = {
    'HKEY_LOCAL_MACHINE\\SOFTWARE\\': [ 'Software.AddReg', ''        ],
    'HKEY_LOCAL_MACHINE\\SYSTEM\\'  : [ 'SetupReg.AddReg', ''        ],
    'HKEY_CURRENT_USER\\'           : [ 'Default.AddReg' , ''        ],
#    'HKEY_USERS\\.DEFAULT\\'        : [ 'Default.AddReg' , ''        ],
    'HKEY_CLASSES_ROOT'             : [ 'Software.AddReg', 'Classes' ]
    }

keyblacklist = [
    'Microsoft\\Cryptography\\RNG', # Random Seed
    'Microsoft\\EventSystem',
    'Microsoft\\Windows\\CurrentVersion\\Explorer\\MountPoints', # Local info about volumes
    'Microsoft\\Windows\\CurrentVersion\\Explorer\\SessionInfo',
    'Microsoft\\Windows\\CurrentVersion\\Explorer\\UserAssist',
    'Microsoft\\Windows\\ShellNoRoam\\MUICache', # Localized string cache
    'Microsoft\\Windows\\CurrentVersion\\Explorer\\MenuOrder',
    'Microsoft\\Windows\\CurrentVersion\\Installer',
    'Microsoft\\Windows\\CurrentVersion\\Uninstall',
    'Microsoft\\Windows\\CurrentVersion\\Internet Settings\\Cache',
    'Microsoft\\Advanced INF Setup',
    'Microsoft\\DirectDraw\\MostRecentApplication',
    'BackupRestore\FilesNotToBackup',# No Backup Files
    '\\ShellNoRoam\\BagMRU',         # Local settings
    'Control\\GroupOrderList',       # Local info
    'Parameters\\Tcpip',             # Tcp/Ip local settings
    'Tcpip\\Parameters',             # Tcp/Ip local settings
    'Dhcp\\Parameters',              # Dchp local settings
    'CurrentControlSet\\',           # This is a link to ControlSet001
    #'\\Control\\DeviceClasses\\',    # Local info
    '\\Enum\\',                      # Local info
    'Classes\\Installer\\Products\\',# Msi
    'LastKnownGoodRecovery'
    ]

# Uncommenting this line disables blacklist, mainly for debug
#keyblacklist = []

fileblacklist = [
    re.compile(r".*?NTUSER.DAT.*?", re.IGNORECASE),
    re.compile(r"\.(pf|log|tmp|sav)$", re.IGNORECASE),
    re.compile(r".*WINDOWS.*Installer.*", re.IGNORECASE),
    re.compile(r".*WINDOWS.*Prefetch.*", re.IGNORECASE),
    re.compile(r".*Temporary Internet Files.*", re.IGNORECASE),
    re.compile(r".*\\Temp\\.*", re.IGNORECASE),
    re.compile(r".*InstallShield Installation Information.*", re.IGNORECASE),
    re.compile(r".*system32\\CatRoot2.*", re.IGNORECASE)
    ]

transtbl = [
    [ re.compile(r"[a-z][:|\?]\\windows.*?", re.IGNORECASE),               '%SystemRoot%' ],
    [ re.compile(r"[a-z][:|\?]\\i386.*?", re.IGNORECASE),                  '%SystemRoot%' ],
    # Programmi - ITALIAN
    # Program Files - ENGLISH
    # more to come...
    [ re.compile(r"^([a-z][:|\?]\\Program(mi|\sfiles)).*?", re.IGNORECASE), '%SystemDrive%\\Programs' ],
    # 8+3 form
    [ re.compile(r"^([a-z][:|\?]\\Progra~\d).*?", re.IGNORECASE),           '%SystemDrive%\\Programs' ],
    [ re.compile(r"\"+([a-z][:|\?]\\Progra[\w|\s]*?)\\.*?", re.IGNORECASE), '%SystemDrive%\\Programs\\\\' ],
    # This should be the last one
    [ re.compile(r"^[a-z][:|\?].*?", re.IGNORECASE),                        '%SystemDrive%']
   ]

unexpandtbl = [
    [ '%SystemDrive%\\Programs',        'Programs',  2 ],
    [ '%SystemRoot%\\System32\\drivers', None,       4 ],
    [ '%SystemRoot%\\System32\\config',  None,       3 ],
    [ '%SystemRoot%\\System32',          None,       2 ],
    [ '%SystemRoot%\\inf',               None,      20 ],
    [ '%SystemRoot%',                    None,       2 ]
    ]


re_comma = re.compile(r'"(.*?)",*')
re_semic = re.compile(r'"(.*?)";*')

suggest = ''

REG_NONE="0x0"
REG_SZ="0x1"
REG_EXPAND_SZ="0x2"
REG_BINARY="0x3"
REG_DWORD="0x4"
REG_MULTI_SZ="0x7"

SECTION=0
PREFIX=1

KEYS=0
VALUES=1
HEADERS=2
FILES=3
FOLDERS=4

WRAP=25
#WRAP=1000
AUTOOPEN=1
ONLYADDED=0
NOPATH=1
CLEANUP=1

QUOTE='"'
NULL='\x00'
BS='\\'
LF='\n'
CR='\r'

CCS='CurrentControlSet'
CS1='ControlSet001'
CS2='ControlSet002'
CS3='ControlSet003'

replacelist = [
    [ BS+CR+LF , ''  ],
    [ BS+CR    , ''  ],
    [ BS+LF    , ''  ],
    [ BS+BS    , BS  ],
    [ CR       , ''  ],
    [ CCS      , CS1 ],
    [ CS2      , CS1 ],
    [ CS3      , CS1 ]
    ]


def unquote(text):
    if len(text)<2:
        return text
    
    if text[0]==QUOTE:
        text=text[1:]
    if text[len(text)-1]==QUOTE:
        text=text[:-1]
    return text

def quote(text):
    if text.find(' ') != -1:
        text = text.replace('\x00', '""')
    else:
        text = text.replace('\x00', '')
        
    if len(text):
        return QUOTE + text + QUOTE
    else:
        return text


def unexpand(value):
    if len(value)<1:
        return value

    for myre in transtbl:
        if value[0] == '@':
            value = '@' + myre[0].sub(myre[1], value[1:])
        else:
            value = myre[0].sub(myre[1], value)
            
    return value

### Pretty print ;)
def wrap(value):
    value = value.replace(' ', '')
    value = value.lower()
    values = value.split(',')
    if len(values) < WRAP:
        return value
    out = ''
    for i in range(0, len(values), WRAP):
        out+= ',' + BS + LF + '  ' + ','.join(values[i:i + WRAP])
    return out[1:] # remove first ,

### decode REG_MULTI_SZ and transform in a list of strings
def decode_sz(value):
    if len(value) < 1:
        return value
    value = value.replace(' ', '')
    value = value.split(',')

    try:
        klist = [ chr(eval('0x' + x)) for x in value ]
    except:
        return ''

    klist = ''.join(klist)
    klist = klist.split('\x00')
    res = []
    for s in klist:
        if len(s):
            res.append(quote(s))
    return ','.join(res)

### Transform a line into a PEBuilder format
def parse_value(Type, value):
    ### REG_DWORD
    if Type == REG_DWORD:
        value = value.replace(' ', '')
        value = value.split(',')
        if len(value) != 4:
            return None
        value.reverse()
        try:
            value = eval('0x' + ''.join(value) + 'L')
        except:
            return None
        return ("0x%x" % long(value))

    ### REG_BINARY
    if Type == REG_BINARY:
        return wrap(value)

    ### REG_MULTI_SZ
    if Type == REG_MULTI_SZ:
        return decode_sz(value)

    for trans in replacelist:
        value = trans[1].join(value.split(trans[0]))
        
    value = unexpand(value)
    return quote(value)

def find_sec(data):
    for key in tkeys.keys():
        if data.find(key) == 0:
            data = data.replace(key, tkeys[key][PREFIX])
            return tkeys[key][SECTION], data
    return None, None

def parse_header(data):
    if len(data) < 2: return None
    name = unquote(data[0].strip())
    value = unquote(data[1].strip())
    return HEADERS, name, value

def parse_files(data):
    what = unquote(data[0]).lower()
    action = unquote(data[1]).lower()
    if action != 'added':
        return None

    if what == 'folders':
        res = FOLDERS
    elif what == 'files':
        res = FILES
    else:
        return None

    data = data[2:]
    fileinfo = None
    if len(data) > 2:
        fileinfo = unquote(data[1])
        
    path = unquote(data[0])

    for f in fileblacklist:
        if f.search(path):
            return None
        if fileinfo and f.search(fileinfo):
            return None

    for myre in transtbl:
        path = myre[0].sub(myre[1], path)
        
    return res, path, fileinfo
    
def parse_reg(data):
    what = unquote(data[0]).lower()
    action = unquote(data[1]).lower()
    data = data[2:]
    
    if action == 'ignored' or action == 'deleted':
        return None
    
    key = unquote(data[0])

    # Blacklisted keys
    for k in keyblacklist:
        if key.lower().find(k.lower()) != -1:
            return None

    # Replace Keys
    for trans in replacelist:
        key = trans[1].join(key.split(trans[0]))

    if what == 'keys' and action == 'added':
        section, keyname = find_sec(key)
        if section is not None:
            return KEYS, section, keyname
        return None

    if what == 'values':
        subkey = unexpand(unquote(data[1]))
        data = data[2:]

        if action == 'changed':
            Type = unquote(data[1])
            if len(data) < 4:
                oldvalue = ''
                value = ''
                Type = 'REG_NONE'
            else:
                oldvalue = unquote(data[2])
                value = unquote(data[3])
        else:
            Type = unquote(data[0])
            if len(data) < 2:
                Type = 'REG_NONE'
                value = ''
            else:
                value = unquote(data[1])

        if ONLYADDED and action == 'changed': return None

        ## Avoid uneeded changed values - FIXME this check should be after processing the value
        if action == 'changed' and value.lower() == oldvalue.lower():
            return None

        section, keyname = find_sec(key)
        if section is not None:
            try:
                Type = eval(Type)
            except:
                return None
            value = parse_value(Type, value)
            if value is None:
                return None

            if subkey == '(Default)':
                subkey = ''

            ### Question: can I do this?
            if (Type == REG_SZ) and (value.find('%') != -1):
                Type = REG_EXPAND_SZ

            kn = keyname.split('\\').pop().lower()
            ### Help dir
            if kn == 'helpdir':
                value = '"%SystemRoot%\Help"'

            if NOPATH:
                try:
                    if (subkey.lower() != 'exepath') and (value.count('\\') == 2):
                        value = value.replace("%SystemRoot%\\system32\\", '')
                except: pass
            return VALUES, action, section, keyname, Type, subkey, value
        return None
        
    stderr.write("Unandled what %s - action %s\n" % (what, action))
    stderr.flush()
    return None

def parse_line(stuff):
    if len(stuff) == 0: return None

    Type = unquote(stuff[0]).lower()
    stuff = stuff[1:]

    if len(Type) < 2:
        return None

    if Type[0] == '-':
        return None

    if Type == 'thing':
        return None
    
    ### Registry entries
    if Type == 'registry':
        return parse_reg(stuff)

    ### Header
    if Type == 'header':
        return parse_header(stuff)

    ### Not needed
    if (Type == 'ini file') or (Type == 'text file') or (Type == 'type'):
        return None

    ### Files
    if Type == 'disk contents':
        return parse_files(stuff)
    
    stderr.write("Unandled Type %s|%s\n" % (Type, stuff))
    stderr.flush()
    return None

def parse_csv(filename):
    inf = { 'registry': { 'Software.AddReg': {}, 'SetupReg.AddReg': {}, 'Default.AddReg' : {} },
            'headers' : {},
            'files'   : {}
            }
            
    lines = open(filename).readlines()

    # Check if separator is , or ;
    first = lines[0]
    lines = lines[1:]
    if len(re_semic.findall(first)):
        sep = re_semic
    else:
        sep = re_comma
            
    for line in lines:
        ### Workaround for quotes in value
        line = line.strip()
        line = line.replace(',"""', ',"\x00')
        line = line.replace('""', '\x00')
        items = sep.findall(line)
        res = parse_line(items)
        if res is not None:
            what = res[0]
            res = res[1:]
            if what == KEYS:
                section = res[0]
                keyname = res[1]
                inf['registry'][section][keyname] = []
            elif what == VALUES:
                res = res[1:]
                section = res[0]
                keyname = res[1]
                Type = res[2]
                subkey = res[3]
                value = res[4]
                
                if not inf['registry'][section].has_key(keyname):
                    inf['registry'][section][keyname] = []

                inf['registry'][section][keyname].append([Type, quote(keyname), quote(subkey), value])
            elif what == HEADERS:
                name = res[0]
                value = res[1]
                inf['headers'][name] = value
            elif what == FILES:
                path = res[0]
                filename = res[1]
                if filename is not None:
                    if inf['files'].has_key(path):
                        inf['files'][path].append(filename)
                    else:
                        inf['files'][path] = [filename]
                else:
                    inf['files'][path] = []
            elif what == FOLDERS:
                if not inf['files'].has_key(res[0]):
                    inf['files'][res[0]] = []
    return inf


def sort_files(filedata):
    dirlist = filedata.keys()
    dirlist.sort()
    start = 0
    res = {}
    for directory in dirlist:
        filelist = filedata[directory]

        if len(filelist) == 0:
            continue

        if directory[-1] == '\\':
            directory = directory[:-1]
        
        letter = letters[start % len(letters)]

        for u in unexpandtbl:
            if directory.lower().find(u[0].lower()) == 0:
                if u[1] is None:
                    res[directory] = [ None, u[2], filelist ]
                else:
                    directory = u[1] + '\\' + directory[len(u[0]) + 1:]
                    if not res.has_key(directory):
                        res[directory] = [ letter, u[2], filelist ]
                        start += 1
                break
    return res
       

def pe_format(inf):
    global suggest
    out = '; Automatic generated by inctrl2pe.py v' + __version__ + LF
    out += ';' + LF
    for k in inf['headers'].keys():
        out += '; %s: %s\n' % (k, inf['headers'][k])
        if k == 'Install program':
            try:
                suggest = inf['headers'][k].split('\\').pop().strip()
                suggest = suggest.split().pop().strip()
            except: pass
    out += LF
    files = ''
    if len(inf['files']):
        dirs = sort_files(inf['files'])
        dirlist = dirs.keys()
        dirlist.sort()

        winntdir = ''

        for i in dirlist:
            if dirs[i][0] is not None:
                winntdir += ("%s=%s,%s" % (dirs[i][0], i, dirs[i][1])) + LF
                destination = dirs[i][0]
            else:
                destination = dirs[i][1]

            for f in dirs[i][2]:
                files += ("files\\%s\\%s=%s,,1" % (i.split('\\',1).pop(), f, destination)) + LF

        if len(winntdir):
            out += LF + "[WinntDirectories]" + LF + winntdir
            
        out += LF
        out += "[SourceDisksFiles]" + LF
        out += files + LF

    for k in inf['registry'].keys():
        if len(inf['registry'][k]):
            out += '['+k+']' + LF
            keys = inf['registry'][k].keys()
            keys.sort()

            if CLEANUP:
                ### cleanup - FIXME doesn't work always
                for subkey in keys[:]:
                    spath = subkey.split('\\')
                    while len(spath):
                        spath = spath[:-1]
                        prepath = '\\'.join(spath)
                        prekey = inf['registry'][k].get(prepath, None)
                        if prekey is not None and len(prekey) == 0:
                            del inf['registry'][k][prepath]
                            keys.remove(prepath)
                
            for subkey in keys:
                if len(inf['registry'][k][subkey]) == 0:
                    skey = quote(subkey)
                    if skey.count(','):
                        out += '0x1,' + skey + ',' + LF
                    else:
                        out += '0x0,' + skey + LF
                    
                for line in inf['registry'][k][subkey]:
                    if line[0] == REG_NONE:
                        while line[-1] == '':
                            line.pop()
                    out += ','.join(line) + LF
                    
            out += LF
    return out            

### Linux/Windows cmdline interface
def cmdline(args):
    if len(args) < 2:
        print 'Usage: %s inputfile.csv' % args[0]
        sys_exit()

    data = parse_csv(args[1])
    inf = pe_format(data)
    print inf
    sys_exit()
    

### Pseudo gui only on win32 + ActiveState python
def gui(args):
    global suggest
    from win32ui import CreateFileDialog
    from win32con import OFN_HIDEREADONLY, OFN_FILEMUSTEXIST, OFN_OVERWRITEPROMPT, OFN_PATHMUSTEXIST
    import os.path

    if len(args) < 2:
        dialog = CreateFileDialog(1, "csv" , None, OFN_HIDEREADONLY | OFN_FILEMUSTEXIST , "Inctrl5 csv Files (*.csv)|*.csv|All Files (*.*)|*.*|", None )
        res = dialog.DoModal()
        if res != 1:
            sys_exit()

        file_in = os.path.join(dialog.GetPathName(), dialog.GetPathName())
    else:
        file_in = args[1]

    data = parse_csv(file_in)
    inf = pe_format(data)

    if len(suggest):
        filename = os.path.splitext(os.path.basename(suggest))[0] + '.inf'
    else:
        filename = os.path.splitext(os.path.basename(file_in))[0] + '.inf'
    
    dialog = CreateFileDialog(0, "inf" , filename, OFN_OVERWRITEPROMPT | OFN_PATHMUSTEXIST, "PE Builder inf (*.inf)|*.inf|All Files (*.*)|*.*|", None )
    res = dialog.DoModal()
    if res != 1:
        sys_exit()

    file_out = os.path.join(dialog.GetPathName(), dialog.GetPathName())
    del dialog  
    open(file_out, 'w').write(inf)        
    if AUTOOPEN:
        import win32api
        win32api.ShellExecute(0, "open", file_out, None, None, 1);

if __name__ == '__main__':
    if platform != 'win32':
        cmdline(argv)
    gui(argv)