#!/usr/bin/env python
# -*- Mode: Python; tab-width: 4 -*-
#
# reg2pe.py - Windows reg file 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.9'

# Changes form 0.8 to 0.9
# - Added AUTOOPEN feature
# Changes from 0.7 to 0.8
# - Added a strip() on each block, since it was missing
#   some handy modified reg files
# - "" now returns null string
# - ControlSet002 and ControlSet003 -> ControlSet001
# Changes from 0.6b to 0.7
# - Added smart path conversion
# - Added hex(7) conversion to string list
# - 'CurrentControlSet' became 'ControlSet001', avoid bsod on startup
# Changes from 0.6 to 0.6b
# - fixed HKEY_USERS key conversion
# Changes from 0.5 to 0.6
# - support for non unicode hex(n)
# - removed null chars at the end of some hex(n)

from codecs import utf_16_decode
from sys import argv, exit, platform
import re

AUTOOPEN=1

### 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' ]
    }

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' ],
        # This should be the last one
        [ re.compile(r"^[a-z][:|\?].*?", re.IGNORECASE),                        '%SystemDrive%']
        ]

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

SECTION=0
PREFIX=1

WRAP=25

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 text[0]==QUOTE:
        text=text[1:]
    if text[len(text)-1]==QUOTE:
        text=text[:-1]
    return text

def quote(text):
    if len(text):
        return QUOTE + text + QUOTE
    else:
        return text
                    
def decode(data):
    data = data.replace(' ', '').split(',')
    try:
        data = [ chr(eval("0x"+x)) for x in data ]
    except:
        return "BAD DATA FORMAT HERE"
    
    text = ''.join(data)
    try:
        text = utf_16_decode(text)[0]
    except:
        return text.replace(NULL, '')

    return text.encode('ascii', 'ignore').replace(NULL, '')


### 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:
        list = [ chr(eval('0x' + x)) for x in value ]
    except:
        return ''

    list = ''.join(list)
    list = list.split('\x00\x00')
    res = []
    for s in list:
        if len(s)>1:
            s = s.replace('\x00','')
            s = s.replace('\\\\', '\\')
            res.append(quote(unexpand(s)))
    return ','.join(res)

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
                
### Returns a string with ascii regfile and \ + \r\n removed
def make_reg(filename):
    ### Very Very Very Nasty hack - The real Horse Power
    ### "Downsample utf16 to ascii"
    stream = open(filename, 'rb').read()
    if stream[:2] == '\xff\xfe' or stream[:2] == '\xfe\xff':
        stream = utf_16_decode(stream)[0]
        stream = stream.encode('ascii', 'ignore')

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

### Parses a section with values
def parse_section(data):
    section = []
    lines = data.split(LF)
    for line in lines:
        if line.find("=")==-1: continue
        key, value = line.strip().split('=',1)
        key = unquote(key)
        value = unquote(value)
        section.append([key, value])
    return section

### Splits reg file in sections
def parse_data(data):
    reg = {}
    reglist = data.split(LF+LF)
    for spl in reglist:
        spl = spl.strip()
        if (len(spl) < 1) or (spl[0]!= '['): continue        
        if spl.count(LF)>0:
            startdata = spl.find(']'+LF)
            name = spl[1:startdata]
            spl=spl[startdata+2:]
            reg[name] = parse_section(spl)
        else:
            reg[spl[1:-1]] = []
    return reg
    
### Wrapping hex values to make nice output
def wrap(value):
    value = value.replace(' ', '')
    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 ,

### Transform a line into a PEBuilder format
def parse_value(value):
    ### dword - REG_DWORD
    if value.find('dword:')==0:
        try:
            value = value.split('dword:').pop()
            value = eval("0x"+value)
        except:
            value = 0
        return REG_DWORD, ("0x%x" % value)

    ### hex - REG_BINARY
    if value.find('hex:')==0:
        return REG_BINARY, wrap(value.split('hex:', 1).pop())

    ### hex(7) - REG_MULTI_SZ
    if value.find('hex(7')==0:
        return REG_MULTI_SZ, decode_sz(value.split('):', 1).pop())

    ### hex(x) - hex string type x
    if value.find('hex(')==0:
        type, value = value.split('):', 1)
        type = "0x" + type[-1:]
        if (type == REG_SZ) or (type == REG_EXPAND_SZ):
            return type, quote(decode(value))
            
        return type, wrap(value)

    value = unexpand(value)
    
    ### string + env - REG_EXPAND_SZ
    if value.find('%')!=-1:
        return REG_EXPAND_SZ, quote(value)

    ### string - REG_SZ
    return REG_SZ, quote(value)

### Convert all the stuff in PEBuilder format then return all file as string
def convert(section, values): 
    section = quote(section)
    out = ''
    if len(values) == 0:
        out += ','.join([REG_NONE, section]) + LF
    for block in values:
        key, value = block
        # @ = (Default)
        if key == "@":
            key = ""
        key = quote(key)
        type, value = parse_value(value)
        out+=','.join([type,section,key,value]) + LF
    return out

### Gets all supported keys from the reg file
def parse_reg(data):
    inf = {
        'Software.AddReg': [],
        'SetupReg.AddReg': [],
        'Default.AddReg' : []
        }
    sections = data.keys()
    sections.sort()
    for section in sections:
        for key in tkeys.keys():
            if section.find(key)==0:
                inf[tkeys[key][SECTION]].append(convert(section.replace(key,tkeys[key][PREFIX]), data[section]))

    return inf

### Master Converter
def reg2pe(inf):
    out = '; Automatic generated by reg2pe.py' + LF
    out += ';' + LF
    for k in inf.keys():
        if len(inf[k]):
            out+= '['+k+']' + LF
            subkeys = inf[k]
            subkeys.sort()
            for subkey in subkeys:
                out+=subkey
    return out

### Linux/Windows cmdline interface
def cmdline(argv):
    if len(argv) < 2:
        print "Usage: %s inputfile.reg" % argv[0]
        exit()
        
    raw = make_reg(argv[1])
    reg = parse_data(raw)
    inf = parse_reg(reg)
    print reg2pe(inf)
    exit()

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

    if len(argv) < 2:
        dialog = CreateFileDialog(1, "reg" , None, OFN_HIDEREADONLY | OFN_FILEMUSTEXIST , "Regedit Files (*.reg)|*.reg|All Files (*.*)|*.*|", None )
        res = dialog.DoModal()
        if res != 1:
            exit()

        file_in = os.path.join(dialog.GetPathName(), dialog.GetPathName())
        filename = os.path.splitext(dialog.GetFileName())[0] + '.inf' 
    else:
        file_in = argv[1]
        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:
        exit()

    file_out = os.path.join(dialog.GetPathName(), dialog.GetPathName())
    del dialog
    raw = make_reg(file_in)
    reg = parse_reg(parse_data(raw))
    inf = reg2pe(reg)
    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)