acquire_mem.py

#! /usr/bin/env python

## Tim R ## 2014 ##
## Multithreaded acquisition of process memory from $pid - Proof of Compulsiv..eeh concept (:
#################
##
## License:
##  It's just some testing code, Use it as breakfast or whatever other purposes you can imagine!
##
#################
##
## Why bother?
##  I had a discussion with someone about the most optimal way
##  to acquire memory from a running process without kernel modules etc. 
##  
## Why multithreaded?
##  I ran into an application which didnt like ptrace hooks and died on me!!
##  To know what her last memories were before she passes away i started some testing and concluded that;
##  a regular script that just iterates over all readable offsets is too slow when racing..
##  
##  This script is slower in use because of the initialization of the main threads and reading 
##  available offsets up-front and re-arranging the threaded output afterwards.. but!
##
##  the ptrace hook happens at the latest possible moment followed immediately by a multithreaded read on all
##  readable offsets, and thats the only place where i wanted to increase performance
##  see the difference for yourself:
##
##  $ time  ./thisscript 3289 |tail -2
##  AAAAAA
##  0.000221967697144
##  real	0m3.080s
##  user	0m2.008s
##  sys	0m0.068s
## 
##  $ time ./regulariterate 3289 |tail -2
##  AAAAAA
##  0.0389778614044
##  real	0m0.076s
##  user	0m0.012s
##  sys	        0m0.056s
##
## Win! Its slower but faster.. rrright?! :)
##
## Bonus?
##  If you try to capture output from a non-daemonized process that just runs for a second or so
##  you can try this wrapper method from a screen or whatever:
##
##  $ cat > catch << EOF
##  #!/bin/bash
##  # if pids found..
##  if [ -n "$(ps x|grep -i process|grep -v grep)" ]; then
##	pids=$(ps x|grep -i process|grep -v 'grep\|pleasegoawaydamnprocess'|awk '{print $1}')
##	#kill -SIGSTOP $pids; echo  "sent SIGSTOP to, $pids"  ## Slow the process down 
##	/usr/bin/python /home/wewqerwer/acquire_mem.py $pids > dump
##	echo "Sleeping 5.....";sleep 5
##	#kill -SIGCONT $pids
##	exit 0
##  fi
## EOF
## $ chmod +x catch ;screen
## Screen version 4.01.00devel (GNU)
## $ while [ 0 ] ; do ./catch ; done
## [detached from 10396.pts-10...]
## $ ./process 
## etcetcetc...
## The SIGSTOP is not ideal but can be used if the application exits before the script finished dumping the contents


import ctypes, re, sys, Queue, threading, time


exitFlag = 0 

class myThread (threading.Thread):
    def __init__(self, threadID, name, q): 
        threading.Thread.__init__(self)
        self.threadID = threadID
        self.name = name
        self.q = q 
    def run(self):
        process_data(self.name, self.q)

def process_data(threadName, q): 
    while not exitFlag:
        queueLock.acquire()
        if not workQueue.empty():
            data = q.get()
            mem_file.seek(data[0])
            chunk = mem_file.read(data[1] - data[0])
            queueLock.release()
            print chunk
            #print "r0: %i r1: %i r2: %s \n" % (data[0], data[1], data[2]),
        else:
            queueLock.release()
        time.sleep(1)

#Interface for PTRACE
c_ptrace = ctypes.CDLL("libc.so.6").ptrace
c_pid_t = ctypes.c_int32 # This assumes pid_t is int32_t
c_ptrace.argtypes = [ctypes.c_int, c_pid_t, ctypes.c_void_p, ctypes.c_void_p]
def ptrace(attach, pid):
    op = ctypes.c_int(16 if attach else 17) #PTRACE_ATTACH or PTRACE_DETACH
    c_pid = c_pid_t(pid)
    null = ctypes.c_void_p()
    err = c_ptrace(op, c_pid, null, null)
    if err != 0: raise SysError, 'ptrace', err

## Parse a line in /proc/$pid/maps.
## Return boundaries of chunk and the read permission character.
def maps_line_range(line):
    m = re.match(r'([0-9A-Fa-f]+)-([0-9A-Fa-f]+) ([-r])', line)
    return [int(m.group(1), 16), int(m.group(2), 16), m.group(3)]

## Dump the readable chunks of memory mapped by a process
if __name__ == "__main__":
    for pid in sys.argv[1:]:

        ## Read memory map from $pid 
        maps_file = open("/proc/" + pid + "/maps", 'r')
        ranges = map(maps_line_range, maps_file.readlines())
        maps_file.close()
        
        ## Initialize threads
        theadList = range(len(ranges)/2)
        queueLock = threading.Lock()
        workQueue = Queue.Queue(0)
        threads = []
        threadID = 0

        for tName in theadList: 
            thread = myThread(threadID, tName, workQueue)
            thread.start()
            threads.append(thread)
            threadID += 1

        t1 = time.time()        
        ## we need to ptrace(PTRACE_ATTACH, $pid) to read /proc/$pid/mem
        ptrace(True, int(pid))
        ## Read the readable mapped ranges
        

        mem_file = open("/proc/" + pid + "/mem", 'r', 0)
        ## Lock queue
        queueLock.acquire()

        ## Spawn all tasks to workers, then release the queue lock
        for r in ranges:
            if r[2] == 'r':
                workQueue.put(r[:])
        queueLock.release()
        t2 = time.time()
 

        ## Empty queue
        while not workQueue.empty():
            pass
 
        exitFlag = 1

        ## Wait for all threads to complete and parse output.
        for t in threads:
            t.join()

        ##Cleanup
        mem_file.close()
        ptrace(False, int(pid))