add patchkit run

Author: lunixbochs

core/patcher.py

@@ -7,14 +7,15 @@
 from binary import Binary
 
 class Patcher:
-    def __init__(self, binary, verbose=False, cflags=None):
+    def __init__(self, binary, verbose=False, cflags=None, silent=False):
         self.bin = Binary(binary)
         self.bin.verbose = verbose
         self.bin.linker.cflags = cflags
         self.patches = []
         self.patchfiles = []
         self.verbose = verbose
         self.cflags = cflags
+        self.silent = silent
 
     def add(self, path):
         if path.endswith('.py'):
@@ -26,12 +27,16 @@ def add(self, path):
                     continue
                 self.patchfiles.append((name, os.path.join(base, os.path.basename(name))))
 
+    def debug(self, *args):
+        if not self.silent:
+            print >>sys.stderr, ' '.join(map(str, args))
+
     def patch(self):
         cwd = os.getcwd()
         try:
             for path, pathname in self.patchfiles:
                 sys.path.insert(0, os.path.dirname(path))
-                print '[*]', pathname
+                self.debug('[*]', pathname)
                 patchfile = os.path.basename(path).rsplit('.', 1)[0]
                 patch = __import__(patchfile)
                 sys.path.pop(0)
@@ -48,8 +53,8 @@ def patch(self):
                             except AttributeError:
                                 pass
                 except Exception:
-                    print 'Warning: could not preserve patch function order'
-                    traceback.print_exc()
+                    self.debug('Warning: could not preserve patch function order')
+                    self.debug(traceback.format_exc())
                     order = vars(patch).values()
 
                 for func in order:
@@ -58,19 +63,19 @@ def patch(self):
                         continue
 
                     if hasattr(func, '__call__'):
-                        print ' [+] %s()' % func.__name__
+                        self.debug(' [+] %s()' % func.__name__)
                         with self.bin.collect() as patchset:
                             try:
                                 func(patchset)
                             except Exception as e:
-                                print 'Exception thrown by patch:', path, func.__name__
+                                self.debug('Exception thrown by patch:', path, func.__name__)
                                 traceback.print_exc()
-                                print 'Memory maps:'
+                                self.debug('Memory maps:')
                                 for prog in self.bin.elf.progs:
                                     if prog.isload:
-                                        print '0x%x-0x%x' % (prog.vaddr, prog.vaddr + prog.vsize)
+                                        self.debug('0x%x-0x%x' % (prog.vaddr, prog.vaddr + prog.vsize))
                                 sys.exit(1)
-                print
+                self.debug()
         finally:
             os.chdir(cwd)
 

run

@@ -0,0 +1,25 @@
+#!/usr/bin/env python
+
+import os
+import sys
+
+from core import Patcher
+from optparse import OptionParser
+
+if __name__ == '__main__':
+    parser = OptionParser(usage='Usage: run [options] <binary> <patchdir> [patchdir...]')
+    parser.add_option('-v', '--verbose', dest="verbose", action="store_true", help="verbose output")
+
+    options, args = parser.parse_args()
+
+    if len(args) < 2:
+        parser.print_help()
+        sys.exit(1)
+
+    args = map(os.path.abspath, args)
+    patchdirs = args[1:]
+
+    patch = Patcher(args[0], verbose=options.verbose, silent=not options.verbose)
+    for d in patchdirs:
+        patch.add(d)
+    patch.patch()