├── .gitignore
├── Module.py
├── README.md
├── dislib.py
└── search.py


/.gitignore:
--------------------------------------------------------------------------------
1 | *.pyc
2 | *.sw[a-z]
3 | 


--------------------------------------------------------------------------------
/Module.py:
--------------------------------------------------------------------------------
  1 | from capstone import *
  2 | import bintools.elf
  3 | import struct
  4 | import dislib
  5 | import re
  6 | 
  7 | class Segment:
  8 |     def __init__(self, name, base, code):
  9 |         self.name = name
 10 |         self.base = base
 11 |         self.code = code
 12 | 
 13 | class Gadget:
 14 |     def __init__(self, byte_sequence, instructions, offset, segment, padding=0):
 15 |         self.byte_sequence = byte_sequence
 16 |         self.instructions = instructions
 17 |         self.base = segment.base + offset
 18 |         self.padding = padding
 19 |         self.segment_name = segment.name
 20 | 
 21 |     def __str__(self):
 22 |         return "%s%s: %s" % ("%s@"%self.segment_name, hex(self.base), '; '.join(self.instructions))
 23 | 
 24 | class Module:
 25 |     def __init__(self, path, architecture):
 26 |         self.name = path.split("/")[-1]
 27 |         self.segments = []
 28 |         self.text_segment = None
 29 |         self.architecture = architecture
 30 |         if self.architecture.lower() == "x86":
 31 |             self.disassembler = Cs(CS_ARCH_X86, CS_MODE_32)
 32 |         elif self.architecture.lower() == "x86_64":
 33 |             self.disassembler = Cs(CS_ARCH_X86, CS_MODE_64)
 34 |         elif self.architecture.lower() == "x86_16": # lol
 35 |             self.disassembler = Cs(CS_ARCH_X86, CS_MODE_16)
 36 |         else:
 37 |             raise Exception("I donut know that architecty")
 38 |         self.load(path)
 39 | 
 40 |     def load(self, path):
 41 |         raise("NO")
 42 | 
 43 |     def find_all_bytes(self, segment, byte_sequence):
 44 |         for i in re.finditer(byte_sequence, segment.code):
 45 |             if i.group(0)[-1] == "\xc3":
 46 |                 padding = 0
 47 |             else:
 48 |                 padding = struct.unpack("H", i.group(0)[-2:])[0]
 49 | 
 50 |             yield self.build_gadget(i.group(0), i.start(), segment)
 51 | 
 52 |     def find_all_returns(self, segment, offset=True):
 53 |         if offset:
 54 |             return self.find_all_bytes(segment, '(?:\xc3|\xc2[\x00-\xFF]{2}){1}')
 55 |         return self.find_all_bytes(segment, '\xc3')
 56 | 
 57 |     def find_all_gadgets(self, lookback, offset=False, remove_dupes=True):
 58 |         start_addresses = []
 59 |         sequences = []
 60 | 
 61 |         for segment in self.segments:
 62 |             for gadget in self.find_all_returns(segment, offset):
 63 |                 gadget_location = gadget.base-segment.base
 64 |                 gadget_length = len(gadget.byte_sequence)
 65 |                 # start at 1 instead of 0 or else you get all rets by default
 66 |                 for offset in range(1, lookback*15):
 67 |                     # don't even look at gadgets that have been examined already
 68 |                     if gadget_location-offset in start_addresses or gadget_location-offset < 0:
 69 |                         continue
 70 | 
 71 |                     byte_sequence = segment.code[gadget_location-offset:gadget_location+gadget_length]
 72 |                     temp_gadget = self.build_gadget(byte_sequence, gadget_location-offset, segment)
 73 | 
 74 |                     if remove_dupes:
 75 |                         if str(temp_gadget).split(": ")[1] in sequences:
 76 |                             continue
 77 | 
 78 |                         sequences.append(str(temp_gadget).split(": ")[1])
 79 | 
 80 |                     if len(temp_gadget.instructions) == lookback:
 81 |                         # if there is a ret in the middle of the gadget, abort
 82 |                         abort = False
 83 |                         for index, value in enumerate(temp_gadget.instructions):
 84 |                             if "ret" in value and index != len(temp_gadget.instructions)-1:
 85 |                                 abort = True
 86 |                                 break
 87 | 
 88 |                         if abort:
 89 |                             continue
 90 | 
 91 |                         # sometimes the last instruction is not actually a return
 92 |                         if "ret" not in temp_gadget.instructions[-1]:
 93 |                             continue
 94 |                         yield temp_gadget
 95 | 
 96 |     def build_gadget(self, byte_sequence, offset, segment):
 97 |         instructions = []
 98 |         for instruction in self.disassembler.disasm(byte_sequence, segment.base+offset):
 99 |             instructions.append("%s %s"%(instruction.mnemonic, instruction.op_str))
100 | 
101 |         return Gadget(byte_sequence, instructions, offset, segment)
102 | 
103 | class PE(Module):
104 |     def load(self, path):
105 |         self.binary = dislib.PEFile(path)
106 |         self.entry_point = entry_point = self.binary.ImageBase + self.binary.EntryPoint
107 |         text_section = self.binary.GetSectionByVA(self.binary.EntryPoint)
108 |         code = text_section.Data[self.binary.EntryPoint - text_section.VA:]
109 |         self.text_segment = Segment(".text", self.entry_point, code)
110 |         self.segments.append(self.text_segment)
111 | 
112 | class ELF(Module):
113 |     def load(self, path):
114 |         import sys
115 |         self.binary = bintools.elf.ELF(path)
116 |         self.imagebase = self.binary.header.entry
117 |         self.entry_point = self.binary.sect_dict['.text'].addr
118 |         self.text_segment = Segment("text", self.entry_point, self.binary.sect_dict['.text'].data)
119 |         for segment in self.binary.sect_dict:
120 |             current_segment = self.binary.sect_dict[segment]
121 |             if current_segment.is_loadable() and current_segment.is_execinstr():
122 |                 self.segments.append(Segment(segment, current_segment.addr, current_segment.data))
123 |                 if segment == '.text':
124 |                     self.text_segment = Segment(segment, current_segment.addr, current_segment.data)
125 | 
126 | class Raw(Module):
127 |     def load(self, path):
128 |         self.binary = None
129 |         self.entry_point = 0
130 |         self.text_segment = Segment(".text", self.entry_point, path)
131 |         self.segments.append(self.text_segment)
132 | 


--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
 1 | #Catfish
 2 | ##Introduction
 3 | Catfish is a tool used ease the process of finding ROP gadgets and creating payloads with them. It is still under development.
 4 | 
 5 | ##Usage
 6 | Right now, Catfish is somewhat limited. Here is a simple demo of the interactive version that will call MessageBoxA using code from SwDir.dll.
 7 | 
 8 |     > load test_bins/swdir.dll
 9 |     > static 0
10 |     > static 0x69218BA0
11 |     > static 0x69218C74
12 |     > static 0
13 |     > call MessageBoxA
14 |     > write messagebox.bin
15 |     >
16 | 
17 | messagebox.bin should now contain a basic payload for launching a MessageBox.
18 | 
19 | ##Prerequisites
20 | All that is needed to run is Distorm (http://code.google.com/p/distorm/)
21 | 
22 | ##Future
23 | Some things planned for the future:
24 | 
25 | * Make it actually work for more than just simple payloads
26 | * Better support for static values on the stack
27 | * 64-bit support
28 | * Use of instruction decomposition
29 | * Automagic ROP chain generation
30 | 


--------------------------------------------------------------------------------
/dislib.py:
--------------------------------------------------------------------------------
  1 | #
  2 | # diSlib64 - diStorm PE library
  3 | # Gil Dabah, 2007
  4 | # http://ragestorm.net/distorm/
  5 | #
  6 | # Thanks to Roee Shenberg for the OO.
  7 | # Thanks to TFKyle for the Impors/Exports/Relocs property idea (loaded upon demand).
  8 | # Thanks to Fred Male for the minor Exports bug fix.
  9 | #
 10 | #Copyright (c) 2007, Gil Dabah
 11 | #All rights reserved.
 12 | #Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met:
 13 | #
 14 | #    * Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer.
 15 | #    * Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer in the documentation and/or other materials provided with the distribution.
 16 | #    * Neither the name of the diStorm nor the names of its contributors may be used to endorse or promote products derived from this software without specific prior written permission.
 17 | #
 18 | #THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 19 | 
 20 | 
 21 | import os
 22 | import sys
 23 | import struct
 24 | import cStringIO
 25 | 
 26 | def ReadAsciiz(buf):
 27 |     try:
 28 |         return buf[:buf.index('\x00')]
 29 |     except ValueError:
 30 |         return None
 31 | 
 32 | MZ_HEADER_LENGTH = 0x40
 33 | PE_HEADER_LENGTH = 0x18
 34 | PE_HEADER_OFFSET = 0x3c
 35 | OPTHDR_SIZE_OFFSET = 0x14
 36 | SECTION_SIZE = 0x28
 37 | 
 38 | IMAGE_FILE_MACHINE_I386 = 0x14c
 39 | IMAGE_FILE_MACHINE_AMD64 = 0x8664
 40 | 
 41 | MAGIC_PE32 = 0x10b
 42 | MAGIC_PE32_PLUS = 0x20b
 43 | 
 44 | class PEException(Exception):
 45 |     pass
 46 | 
 47 | def SafeRead(f, length):
 48 |     """Tries to read and will raise an exception if there are not enough bytes."""
 49 |     data = f.read(length)
 50 |     if len(data) != length:
 51 |         raise IOError("Couldn't read enough bytes!")
 52 |     return data
 53 | 
 54 | def GetQword(s, offset, count = 1):
 55 |     if count == 1:
 56 |         return struct.unpack("<Q", s[offset:offset + 8])[0]
 57 |     else:
 58 |         return struct.unpack("<%dQ" % count ,s[offset:offset + 8*count])
 59 | 
 60 | def GetDword(s, offset, count = 1):
 61 |     if count == 1:
 62 |         return struct.unpack("<L", s[offset:offset + 4])[0]
 63 |     else:
 64 |         return struct.unpack("<%dL" % count ,s[offset:offset + 4*count])
 65 | 
 66 | def GetShort(s, offset, count = 1):
 67 |     if count == 1:
 68 |         return struct.unpack("<H", s[offset:offset + 2])[0]
 69 |     else:
 70 |         return struct.unpack("<%dH" % count, s[offset:offset + 2*count])
 71 | 
 72 | class _MZHeader(object):
 73 |     """MZ File Header object. Intended for internal use only."""
 74 |     def __init__(self,f):
 75 |         try:
 76 |             f.seek(0)
 77 |             MZHeader = f.read(MZ_HEADER_LENGTH)
 78 |             if len(MZHeader) < MZ_HEADER_LENGTH:
 79 |                 raise PEException("Cannot read MZ header.")
 80 |         except IOError:
 81 |             raise PEException("Cannot read MZ header.")
 82 | 
 83 |         if (MZHeader[:2] != "MZ"):
 84 |             raise PEException("Invalid MZ file.")
 85 | 
 86 |         self.PEHeaderOffset = GetDword(MZHeader, PE_HEADER_OFFSET)
 87 | 
 88 | class PEFile(object):
 89 |     def __init__(self, FileName, NewImageBase = None):
 90 |         """Filename can be either a file object or filename."""
 91 |         ImportsOff = 0
 92 |         ExportsOff = 0
 93 |         RelocationsOff = 0
 94 |         RelocationsSize = 0
 95 |         self._Imports = None
 96 |         self._Exports = None
 97 |         self._Relocs = None
 98 |         # Holds the delta between the original ImageBase and the new one.
 99 |         self.Delta = 0
100 |         if isinstance(FileName, basestring):
101 |             try:
102 |                 f = open(FileName, "rb")
103 |             except IOError:
104 |                 raise PEException("File not found or cannot be accessed.")
105 |         elif isinstance(FileName, file):
106 |             f = FileName
107 |         else:
108 |             raise PEException("Invalid file/filename argument!")
109 | 
110 |         self._MZHeader = _MZHeader(f)
111 |         # People might want to check this one out.
112 |         self.PEHeaderOffset = self._MZHeader.PEHeaderOffset
113 | 
114 |         f.seek(self.PEHeaderOffset)
115 |         try:
116 |             PEHeader = SafeRead(f, PE_HEADER_LENGTH)
117 |         except IOError:
118 |             raise PEException("Cannot read PE header.")
119 | 
120 |         if (PEHeader[:4] != "PE\0\0"):
121 |             raise PEException("Invalid PE file.")
122 | 
123 |         self.MachineType = GetShort(PEHeader, 4)
124 |         if self.MachineType != IMAGE_FILE_MACHINE_I386 and self.MachineType != IMAGE_FILE_MACHINE_AMD64:
125 |             raise PEException("PE file is not an i386+ image.")
126 | 
127 |         OptHeaderSize = GetShort(PEHeader, OPTHDR_SIZE_OFFSET)
128 |         try:
129 |             #OptHeader = SafeRead(f, OptHeaderSize)
130 |             OptHeader = f.read(0xf8)
131 |         except IOError:
132 |             raise PEException("Cannot read optional PE header.")
133 | 
134 |         self.OptMagic = GetShort(OptHeader, 0)
135 |         if self.OptMagic != MAGIC_PE32 and self.OptMagic != MAGIC_PE32_PLUS:
136 |             raise PEException("Optional header magic isn't not a PE32+.")
137 | 
138 |         CodeSize = GetDword(OptHeader, 4)
139 | 
140 |         EntryPoint, CodeBase = GetDword(OptHeader, 0x10, 2)
141 |         if (self.OptMagic == MAGIC_PE32):
142 |             DataBase, ImageBase = GetDword(OptHeader, 0x18, 2)
143 |         else:
144 |             DataBase = 0 # No DataBase for AMD64
145 |             ImageBase = GetQword(OptHeader, 0x18)
146 | 
147 |         SectionAlignment, FileAlignment = GetDword(OptHeader, 0x20, 2)
148 |         SectionAlignment -= 1
149 |         FileAlignment -= 1
150 | 
151 |         # Delta holds the difference between the original image base and the given one,
152 |         # this will be used in order to apply the relocations correctly.
153 |         if (NewImageBase != None):
154 |             self.Delta = NewImageBase - ImageBase
155 |             ImageBase = NewImageBase
156 | 
157 |         NumOfSections = GetShort(PEHeader, 6)
158 |         try:
159 |             f.seek(self.PEHeaderOffset + OptHeaderSize + 0x18)
160 |             RawSections = SafeRead(f, NumOfSections * SECTION_SIZE)
161 |         except IOError:
162 |             raise PEException("Cannot read in sections.")
163 | 
164 |         DirectoriesOffset = 0x60
165 |         if (self.OptMagic == MAGIC_PE32_PLUS):
166 |             DirectoriesOffset += 0x10
167 |         NumberOfRvaAndSizes = GetDword(OptHeader, DirectoriesOffset-4)
168 |         if 0 < NumberOfRvaAndSizes:
169 |             try:
170 |                 ExportsOff = GetDword(OptHeader, DirectoriesOffset)
171 |             except:
172 |                 pass
173 |         if 1 < NumberOfRvaAndSizes:
174 |             try:
175 |                 ImportsOff = GetDword(OptHeader, DirectoriesOffset + 8)
176 |             except:
177 |                 pass
178 |         if 5 < NumberOfRvaAndSizes:
179 |             try:
180 |                 RelocationsOff = GetDword(OptHeader, DirectoriesOffset + 0x28)
181 |                 RelocationsSize = GetDword(OptHeader, DirectoriesOffset + 0x2c)
182 |             except:
183 |                 pass
184 | 
185 |         try:
186 |             self.SizeOfHeaders = GetDword(OptHeader, 0x3c)
187 |         except:
188 |             raise PEException("Cannot read SizeOfHeaders")
189 |         try:
190 |             self.SizeOfImage = GetDword(OptHeader, 0x38)
191 |         except:
192 |             raise PEException("Cannot read SizeOfImage")
193 | 
194 |         f.seek(0)
195 |         FileMap = f.read()
196 |         f.close()
197 |         Sections = []
198 | 
199 |         for i in xrange(0, NumOfSections*SECTION_SIZE, SECTION_SIZE):
200 |             SecName = RawSections[i:i + 8].strip('\x00')
201 |             SecVA, SecSize, SecFileOffset = GetDword(RawSections, i + 0xc, 3)
202 |             SecFlags = GetDword(RawSections, i + 0x24)
203 |             # Align everything.
204 |             SecVA &= ~SectionAlignment
205 |             if (SecVA < self.SizeOfHeaders):
206 |                 SecVA = 0
207 |             SecFileOffset &= ~(FileAlignment)
208 |             SecSize = (SecSize + FileAlignment) & ~FileAlignment
209 |             try:
210 |                 SecBin = FileMap[SecFileOffset:SecFileOffset + SecSize]
211 |                 # Align the section to SectionAlignment and pad section with zeros.
212 |                 alignSize = SectionAlignment + 1 - (len(SecBin) % (SectionAlignment+1))
213 |                 if (alignSize != SectionAlignment + 1):
214 |                     SecBin += '\x00' * alignSize
215 |             except:
216 |                 raise PEException("Cannot read section or size invalid.")
217 |             Sections.append(Section(i / SECTION_SIZE, SecName, SecVA, SecFlags, SecBin))
218 | 
219 |         # Create a section for the whole file, in case there are VA's which point to no section.
220 |         # Make sure it's the last section so GetSectionByVA scans it lastly.
221 |         # Also align this section to SectionAlignment.
222 |         Sections.append(Section(-1, "FILEMAP", 0, 0, FileMap + '\x00' * (SectionAlignment + 1 - (len(FileMap) % (SectionAlignment+1)))))
223 | 
224 |         self.Sections = Sections
225 |         self.CodeSize = CodeSize
226 |         self.EntryPoint = EntryPoint
227 |         self.CodeBase = CodeBase
228 |         self.DataBase = DataBase
229 |         self.ImageBase = ImageBase
230 |         self.RelocationsSize = RelocationsSize
231 |         self.RelocationsOff = RelocationsOff
232 |         self.ImportsOff = ImportsOff
233 |         self.ExportsOff = ExportsOff
234 | 
235 |     def _GetImports(self):
236 |         """ Getter for self.Imports. """
237 |         if self._Imports != None or not self.ImportsOff:
238 |             return self._Imports
239 |         # Loads and caches imports upon first request.
240 |         self._Imports = self.LoadImports()
241 |         return self._Imports
242 | 
243 |     def _GetExports(self):
244 |         """ Getter for self.Exports. """
245 |         if self._Exports != None or not self.ExportsOff:
246 |             return  self._Exports
247 |         # Loads and caches exports upon first request.
248 |         self._Exports = self.LoadExports()
249 |         return self._Exports
250 | 
251 |     def _GetRelocs(self):
252 |         """ Getter for self.Relocs. """
253 |         if self._Relocs != None or not self.RelocationsOff:
254 |             return self._Relocs
255 |         # Loads and caches relocations and apply them upon first request.
256 |         self._Relocs = self.GetRelocations()
257 |         self.ApplyRelocations()
258 |         return self._Relocs
259 | 
260 |     # Use properties for Imports/Exports/Relocs, so they will be loaded only once upon demand.
261 |     Imports = property(_GetImports)
262 |     Exports = property(_GetExports)
263 |     Relocs = property(_GetRelocs)
264 | 
265 |     def GetSectionByVA(self, VA):
266 |         for i in self.Sections:
267 |             if (VA >= i.VA) and (VA < (i.VA + i.Size)):
268 |                 return i
269 |         return None
270 | 
271 |     def LoadImports(self):
272 |         ImpSec = self.GetSectionByVA(self.ImportsOff)
273 |         VA = self.ImportsOff
274 |         ImageBase = self.ImageBase
275 |         if (ImpSec == None):
276 |             raise PEException("Cannot find import table section.")
277 |         Imports = [] 
278 |         VA -= ImpSec.VA
279 |         while True:
280 |             try:
281 |                 ModuleNameVA = GetDword(ImpSec.Data, VA + 0xc)
282 |             except IndexError:
283 |                 raise PEException("Cannot read imported module name RVA.")
284 |             try:
285 |                 FirstThunkVA = GetDword(ImpSec.Data, VA + 0x10)
286 |             except IndexError:
287 |                 raise PEException("Cannot read imported functions addresses RVA.")
288 |             # Stop parsing import descriptors on any of the next cases.
289 |             if ModuleNameVA == 0 or FirstThunkVA == 0:
290 |                 break
291 |             sec = self.GetSectionByVA(ModuleNameVA)
292 |             if (sec == None):
293 |                 raise PEException("Cannot find imported module name.")
294 | 
295 |             ModuleName = ReadAsciiz(sec.Data[ModuleNameVA - sec.VA:])
296 |             if (ModuleName == None):
297 |                 raise PEException("Cannot read imported module name.")
298 |             FunctionsNamesAddrs = []
299 |             try:
300 |                 OrigFirstThunkVA = GetDword(ImpSec.Data, VA)
301 |             except IndexError:
302 |                 raise PEException("Cannot read imported functions names RVA.")
303 |             PtrsSec = self.GetSectionByVA(FirstThunkVA)
304 |             if (PtrsSec == None):
305 |                 raise PEException("Cannot find imported functions addresses.")
306 | 
307 |             # See if we can use Original Thunk or the First Thunk.
308 |             if ((OrigFirstThunkVA == 0) or (OrigFirstThunkVA < self.SizeOfHeaders) or (OrigFirstThunkVA >= self.SizeOfImage)):
309 |                 # Continue anyways, some linkers use a zero OrigFirstThunkVA...
310 |                 # Prefer using First Thunk.
311 |                 OrigFirstThunkVA = FirstThunkVA
312 | 
313 |             sec = self.GetSectionByVA(OrigFirstThunkVA)
314 |             if (sec == None):
315 |                 raise PEException("Cannot find imported functions names.")
316 | 
317 |             OrigFirstThunkVA -= sec.VA
318 |             # Scan for the imported functions.
319 |             while True:
320 |                 # Get an RVA to its ImportByName structure.
321 |                 try:
322 |                     if (self.OptMagic == MAGIC_PE32):
323 |                         ImportByNameVA = GetDword(sec.Data, OrigFirstThunkVA)
324 |                         OrdBit = 1L << 31
325 |                     else:
326 |                         ImportByNameVA = GetQword(sec.Data, OrigFirstThunkVA)
327 |                         OrdBit = 1L << 63
328 |                 except IndexError:
329 |                     raise PEException("Cannot read ImportByName RVA.")
330 |                 # Is it the end of list?
331 |                 if (ImportByNameVA == 0):
332 |                     break
333 |                 # Is this offset is actually an ordinal number?
334 |                 if (ImportByNameVA & OrdBit):
335 |                     FuncName = "Ord_0x%04X" % (ImportByNameVA & 0xffff)
336 |                 else:
337 |                     # Move next.
338 |                     ImportByNameVA += 2
339 |                     sec2 = self.GetSectionByVA(ImportByNameVA)
340 |                     if (sec2 == None):
341 |                         raise PEException("Cannot find imported function name section.")
342 | 
343 |                     ImportByNameVA -= sec2.VA
344 |                     # Get the imported function name.
345 |                     FuncName = ReadAsciiz(sec2.Data[ImportByNameVA:])
346 |                     if (FuncName == None):
347 |                         raise PEException("Cannot read imported function name.")
348 | 
349 |                     if (len(FuncName) == 0):
350 |                         break
351 |                 # Address of pointer to function in the IAT.
352 |                 FunctionsNamesAddrs.append((FirstThunkVA + ImageBase, FuncName))
353 |                 if (self.OptMagic == MAGIC_PE32):
354 |                     OrigFirstThunkVA += 4
355 |                     FirstThunkVA += 4
356 |                 else:
357 |                     OrigFirstThunkVA += 8
358 |                     FirstThunkVA += 8
359 |             # Create them all as Import class. Add the module and its functions to the list.
360 |             Imports += [Import(va, fname, ModuleName) for va, fname in FunctionsNamesAddrs]
361 |             # Next module.
362 |             VA += 0x14
363 |         return Imports
364 | 
365 |     def LoadExports(self):
366 |         ExpSec = self.GetSectionByVA(self.ExportsOff)
367 |         VA = self.ExportsOff
368 |         CodeBase = self.CodeBase
369 |         ImageBase = self.ImageBase
370 |         if (ExpSec == None):
371 |             raise PEException("Cannot find export table section.")
372 | 
373 |         if (len(ExpSec.Data) < 0x28):
374 |             raise PEException("Export table too small!")
375 | 
376 |         VA -= ExpSec.VA
377 |         # Read basic information about the Export Table.
378 |         OrdinalBase, FuncsCount, NamesCount, FuncsAddrsVA = GetDword(ExpSec.Data, VA + 0x10, 4)
379 | 
380 |         # Get the functions' addresses array page.
381 |         FuncsSec = self.GetSectionByVA(FuncsAddrsVA)
382 |         if (FuncsSec == None):
383 |             raise PEException("Cannot find export table functions addresses.")
384 | 
385 |         FuncsAddrsVA -= FuncsSec.VA
386 |         
387 |         NamesAddrsVA = GetDword(ExpSec.Data, VA + 0x20)
388 |         NamesSec = self.GetSectionByVA(NamesAddrsVA)
389 |         if (NamesSec == None):
390 |             raise PEException("Cannot find export table functions names.")
391 | 
392 |         NamesAddrsVA -= NamesSec.VA
393 | 
394 |         # Get the functions' ordinals array page.
395 |         OrdsAddrsVA = GetDword(ExpSec.Data, VA + 0x24)
396 |         OrdsSec = self.GetSectionByVA(OrdsAddrsVA)
397 |         if (OrdsSec == None):
398 |             raise PEException("Cannot find export table functions ordinals.")
399 | 
400 |         OrdsAddrsVA -= OrdsSec.VA
401 | 
402 |         ordnum = 0
403 |         FuncNames = []
404 |         Names = []
405 |         Ords = []
406 |         # [addr, isUsed]
407 |         if ((FuncsAddrsVA + FuncsCount*4) > len(FuncsSec.Data)):
408 |             raise PEException("Exported functions count exceeds sections.")
409 | 
410 |         # Read in all addresses.
411 |         if (FuncsCount == 1):
412 |             Addrs = [[GetDword(FuncsSec.Data, FuncsAddrsVA, 1), 0]]
413 |         else:
414 |             Addrs = [[i, 0] for i in GetDword(FuncsSec.Data, FuncsAddrsVA, FuncsCount)]
415 | 
416 |         if ((NamesAddrsVA + NamesCount*4) > len(NamesSec.Data)):
417 |             raise PEException("Exported functions count exceeds sections.")
418 | 
419 |         # Read in these functions' names.
420 |         for i in xrange(NamesCount):
421 |             NameVA = GetDword(NamesSec.Data, NamesAddrsVA + i*4)
422 |             NameSec = self.GetSectionByVA(NameVA)
423 |             if (NameSec == None):
424 |                 raise PEException("Cannot find exported function name.")
425 | 
426 |             NameVA -= NameSec.VA
427 |             Name = ReadAsciiz(NameSec.Data[NameVA:])
428 |             if (Name == None):
429 |                 raise PEException("Cannot read exported function name.")
430 | 
431 |             Names.append(Name)
432 | 
433 |         if ((OrdsAddrsVA + NamesCount*2) > len(OrdsSec.Data)):
434 |             raise PEException("Exported functions count exceeds sections.")
435 | 
436 |         # Read the ordinals of these functions.
437 |         if (NamesCount == 1):
438 |             Ords = [GetShort(OrdsSec.Data, OrdsAddrsVA, 1)]
439 |         else:
440 |             Ords = GetShort(OrdsSec.Data, OrdsAddrsVA, NamesCount)
441 | 
442 |         # [name, ord]
443 |         # Mark the functions which have a name and an ordinal.
444 |         for i in enumerate(Names):
445 |             j = Ords[i[0]]
446 |             FuncNames.append(Export(i[1], j + OrdinalBase, Addrs[j][0]))
447 |             # Mark exported function address as used.
448 |             Addrs[j][1] = 1
449 | 
450 |         # Add functions that exported by ordinals only by exclusion.
451 |         # This is a bit tricky, that's why we had a flag indicating whether the function was already ready to go...
452 |         for i in [i for i in enumerate(Addrs) if i[1][1] != 1]:
453 |             j = i[0] + OrdinalBase
454 |             FuncNames.append(Export("", j, Addrs[j][0]))
455 |         return FuncNames
456 | 
457 |     def GetRelocations(self):
458 |         VA = self.RelocationsOff
459 |         Size = self.RelocationsSize
460 |         Sections = self.Sections
461 |         # Read a page header.
462 |         RelocatedAddrs = []
463 |         RelocSec = self.GetSectionByVA(VA)
464 |         if (RelocSec == None):
465 |             raise PEException("Cannot find base relocations section.")
466 | 
467 |         VA -= RelocSec.VA
468 |         Index = 0
469 |         while VA < Size:
470 |             if (VA + 8 > Size):
471 |                 raise PEException("Not enough bytes to read relocations header.")
472 |             # Read a page header.
473 |             PageVA, RelocCount = GetDword(RelocSec.Data, VA, 2)
474 |             Index += 1
475 |             # Calc the number of relocations in this page.
476 |             RelocCount -= 8 #exclude header
477 |             if (VA + RelocCount > Size):
478 |                 raise PEException("Relocations count exceeds section.")
479 | 
480 |             VA += 8
481 |             # Read in the addresses.
482 |             for i in xrange(RelocCount/2):
483 |                 # Is it the end of array? Halt.
484 |                 if (VA >= Size):
485 |                     break
486 |                 TypeOffset = GetShort(RelocSec.Data, VA)
487 |                 # Is it the end of array? Halt.
488 |                 if (TypeOffset == 0):
489 |                     # Skip ABSOLUTE relocation type, move to next page.
490 |                     VA += RelocCount-i*2
491 |                     break
492 |                 VA += 2
493 |                 # Extract the type and offset from the word sized variable.
494 |                 Type = TypeOffset >> 12
495 |                 Offset = TypeOffset & 0xfff
496 |                 # Executables of PE type support only HIGHLOW type.
497 |                 if (Type == 3):
498 |                     PageSec = self.GetSectionByVA(PageVA + Offset)
499 |                     if (PageSec != None):
500 |                         # Append the data into the list as a Relocation class with the original relocated dword as well.
501 |                         RelocatedAddrs.append(Relocation(PageSec, PageVA, Offset, GetDword(PageSec.Data, PageVA + Offset - PageSec.VA)))
502 |                 elif (Type == 10):
503 |                     PageSec = self.GetSectionByVA(PageVA + Offset)
504 |                     if (PageSec != None):
505 |                         # Append the data into the list as a Relocation class with the original relocated qword as well.
506 |                         RelocatedAddrs.append(Relocation(PageSec, PageVA, Offset, GetQword(PageSec.Data, PageVA + Offset - PageSec.VA)))
507 | 
508 |         return RelocatedAddrs
509 | 
510 |     def ApplyRelocations(self):
511 |         # Don't assume VA's are sorted properly.
512 |         # We sort them out, because the algorithm for rebuilding the relocated section assumes
513 |         # the relocation's addresses are from low to high addresses.
514 |         self._Relocs.sort(SortRelocCallback)
515 |         # Hold relocations by their sections.
516 |         Secs = [[] for i in xrange(len(self.Sections))]
517 |         # Sort them into their corresponding section.
518 |         for i in self._Relocs:
519 |             Secs[i.Section.Index].append(i)
520 |         # Eliminate empty sections (which don't have relocations).
521 |         Secs = filter(len, Secs)
522 |         # Now for every section apply its changes
523 |         for Sec in Secs:
524 |             OldVA = 0
525 |             # TmpBin will be used for efficiency in order to apply the patch to Python buffers.
526 |             TmpBin = cStringIO.StringIO()
527 |             # Get its section.
528 |             SecVA = Sec[0].Section.VA
529 |             PageSec = self.GetSectionByVA(SecVA)
530 |             # Patches the required addresses.
531 |             for i in Sec:
532 |                 # Copy the buffer till this address.
533 |                 buf = PageSec.Data[OldVA:i.RelocAddr - SecVA]
534 |                 if len(buf):
535 |                     TmpBin.write(buf)
536 |                 if (self.OptMagic == MAGIC_PE32):
537 |                     # Calc next offset.
538 |                     OldVA = i.RelocAddr - SecVA + 4
539 |                     # Append the new relocated value.
540 |                     TmpBin.write(struct.pack("<L", i.OriginalValue + self.Delta))
541 |                 else:
542 |                     # Calc next offset.
543 |                     OldVA = i.RelocAddr - SecVA + 8
544 |                     # Append the new relocated value.
545 |                     TmpBin.write(struct.pack("<Q", i.OriginalValue + self.Delta))
546 |             # Append any tail buffer left.
547 |             TmpBin.write(PageSec.Data[OldVA:])
548 |             # And use the new section binary.
549 |             PageSec.Data = TmpBin.getvalue()
550 |             TmpBin.close()
551 | 
552 | class Section(object):
553 |     def __init__(self, Index, Name, VA, Flags, Bin):
554 |         # Number of the section in the file.
555 |         self.Index = Index
556 |         self.Name = Name
557 |         # Pointer to File Raw Data.
558 |         self.VA = VA
559 |         self.Flags = Flags
560 |         self.Data = Bin
561 |         self.Size = len(Bin)
562 |     def __repr__(self):
563 |         return "<section %s, VA = 0x%.8X, Flags = 0x%.8X, <0x%X bytes data>>" % (self.Name, self.VA, self.Flags, len(self.Data))
564 | 
565 | class Export(object):
566 |     def __init__(self, fname, ord, va):
567 |         # The VA in which the function sits.
568 |         self.VA = va
569 |         # Function name (if exists).
570 |         if len(fname) == 0 or fname == None:
571 |             self.Name = "<Unnamed>"
572 |         self.Name = fname
573 |         # Function ordinal.
574 |         self.Ordinal = ord
575 |     def __repr__(self):
576 |         return "<Exported Function %s, Ordinal %d, VA 0x%.8X>" % (self.Name, self.Ordinal, self.VA)
577 | 
578 | class Import(object):
579 |     def __init__(self, va, fname, modname):
580 |         self.VA = va
581 |         self.Name = fname
582 |         self.ModuleName = modname
583 |     def __repr__(self):
584 |         return "<Imported Function %s.%s at 0x%.8X>" % (self.ModuleName[:self.ModuleName.rfind('.')], self.Name, self.VA)
585 | 
586 | class Relocation(object):
587 |     def __init__(self, Section, PageVA, Offset, OriginalValue):
588 |         self.Section = Section
589 |         self.PageVA = PageVA
590 |         self.Offset = Offset
591 |         # Original value of the dword which isn't relocated yet.
592 |         self.OriginalValue = OriginalValue
593 |     def _get_reloc_address(self):
594 |         return self.PageVA + self.Offset
595 |     def __repr__(self):
596 |         return "<Relocation at 0x%.8X, OrigVal=%08x>" % (self.RelocAddr, self.OriginalValue)
597 |     RelocAddr = property(_get_reloc_address)
598 | 
599 | def SortRelocCallback(R1, R2):
600 |     " Helper function for sorting all relocations' addresses. "
601 |     res = R1.RelocAddr - R2.RelocAddr
602 |     if res < 0:
603 |         return -1
604 |     elif res > 0:
605 |         return 1
606 |     return 0
607 |     
608 | def main():
609 |     if (len(sys.argv) != 2) and (len(sys.argv) != 3):
610 |         print os.path.split(sys.argv[0])[1] + " <PE file name> <new image base in hex>"
611 |         return None
612 |     filename = sys.argv[1]
613 |     NewImageBase = None
614 |     if (len(sys.argv) == 3):
615 |         NewImageBase = int(sys.argv[2], 16)
616 |     PEObj = PEFile(filename, NewImageBase)
617 |     print "diSlib, http://ragestorm.net/distorm/\n"
618 |     print "Image Base: 0x%08x, Code Size: 0x%x" % (PEObj.ImageBase, PEObj.CodeSize)
619 |     print "Entry Point RVA: %08x" % PEObj.EntryPoint
620 |     print "Sections:"
621 |     for i in PEObj.Sections[:-1]:
622 |         print "%d.Name: %s, VA: %x, Size: %x, Flags: %x" % (i.Index + 1, i.Name, i.VA, i.Size, i.Flags)
623 |     if PEObj.Imports:
624 |         print "Imports:"
625 |         for i in PEObj.Imports:
626 |             print i
627 |     try:
628 |         if PEObj.Exports:
629 |             print "Exports:"
630 |             for i in PEObj.Exports:
631 |                 print i
632 |     except:
633 |         # Ignore corrupted exports.
634 |         pass
635 |     if PEObj.Relocs:
636 |         print "Relocations:"
637 |         for i in PEObj.Relocs:
638 |             print i
639 | 
640 |     try:
641 |         # If diStorm isn't available, we won't disassemble anything.
642 |         import distorm3
643 | 
644 |         DecodeType = distorm3.Decode32Bits
645 |         if (PEObj.MachineType == IMAGE_FILE_MACHINE_AMD64):
646 |             DecodeType = distorm3.Decode64Bits
647 | 
648 |         # Find code section and disassemble entry point routine.
649 |         TextSec = PEObj.GetSectionByVA(PEObj.EntryPoint)
650 |         if TextSec == None:
651 |             return
652 |         l = distorm3.Decode(PEObj.ImageBase + PEObj.EntryPoint, TextSec.Data[PEObj.EntryPoint - TextSec.VA:][:4*1024], DecodeType)
653 |         for i in l:
654 |             print "0x%08x (%02x) %-20s %s" % (i[0],  i[1],  i[3],  i[2])
655 |             if ((i[2][:3] == "RET") or (i[2] == "INT 3") or (i[2][:3] == "JMP")):
656 |                 break
657 |     except:
658 |         pass
659 | 
660 | if (__name__ == "__main__"):
661 |     main()
662 | 


--------------------------------------------------------------------------------
/search.py:
--------------------------------------------------------------------------------
 1 | import Module, sys
 2 | 
 3 | if len(sys.argv) < 2:
 4 |     print "%s [filename] [instructions]"%sys.argv[0]
 5 |     sys.exit(1)
 6 | 
 7 | filename = sys.argv[1]
 8 | lookback = int(sys.argv[2])
 9 | 
10 | if ".exe" in filename:
11 |     m = Module.PE(filename, "x86")
12 | else:
13 |     m = Module.ELF(filename, "x86")
14 | 
15 | for i in m.find_all_gadgets(lookback):
16 |     print i
17 | 


--------------------------------------------------------------------------------