'''
Copyright (c) 2007 Paolo Palumbo
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:
 
* The above copyright notice and this permission notice shall be included in
  all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
THE SOFTWARE.
'''

'''
Version 1.0:
* Created the script

Version 1.1:
* Fix: could not create names if the name was already present in the db
'''

# Analysis part
def main():
    
    gOrdinal = 0                 # for strings issues
    
    # Check the existence of DllFunctionCall function
    AddressOfDllCall = get_name_ea(BADADDR,"DllFunctionCall")
    if AddressOfDllCall == BADADDR:
        print "[DllFunctionCall] Cannot find DllFunctionCall!"
        return
    
    # Get Xrefs to the function
    DllFunctionCallsReferences = CodeRefsTo(AddressOfDllCall,0)
    
    # Iterate through references
    for Reference in DllFunctionCallsReferences:
        
        print "[DllFunctionCall] Use at 0x%x" % Reference
        
        # Move to where we believe the stub starts
        StubStartAddress = Reference - 0x15
        
        # Do check if this is the stub we support
        if ( (Byte(StubStartAddress) == 0xA1) and (Word(StubStartAddress + 5) == 0xC00B) and \
             (Byte(StubStartAddress + 7) == 0x74) and (Word(StubStartAddress + 9) == 0xE0FF) ):
            print "[DllFunctionCall] Stub start at 0x%x" % StubStartAddress
            
            # Create function if IDA did not create it already
            if not get_func(Reference):
                print "[DllFunctionCall] Creating a function for the stub"
                MakeFunction(StubStartAddress,BADADDR)
                
            # Retrieve address of the DllFunctionCallStructure
            OffsetToDllFunctionCallStructure = GetOperandValue(StubStartAddress + 0xB,0)
            
            # Read the offset to wrapped function
            OffsetToDllFunctionName = Dword(OffsetToDllFunctionCallStructure + 4)
            
            # Convert that to string
            MakeUnkn(OffsetToDllFunctionName,3)
            make_ascii_string(OffsetToDllFunctionName,0x0,ASCSTR_C)
            
            # Read the name and name the current properly the current stub. Note that IDA will
            # add type info where available :)
            FunctionName = get_ascii_contents(OffsetToDllFunctionName,get_max_ascii_length(OffsetToDllFunctionName,ASCSTR_C),ASCSTR_C)
            if(MakeName(StubStartAddress,FunctionName) == 0):
                # We cannot create the string, create alternate string
                gOrdinal += 1
                MakeName(StubStartAddress,FunctionName + "_%i" % gOrdinal) 
            
            MakeName(Dword(StubStartAddress+1),'g'+FunctionName)
            print "[DllFunctionCall] Stub successfully created!"
        else:
            print "[DllFunctionCall] Unsupported DllFunctionCall use at 0x%x" % Reference


# Main function
print "[DllFunctionCall] DllFunctionCall by Paolo Palumbo"
print "[DllFunctionCall] Fixing DllFunctionCall stubs..."
main()
print "[DllFunctionCall] Thank you for using me!! Have a nice day!"
print "[DllFunctionCall] Feedback & Bug reports at paolo.palumbo@hotmail.com"