adding arch checking to memory_exec issue #40

.gitignore

@@ -30,7 +30,7 @@ client/sources/resources_msvcr90_dll.c
 __pycache__/
 client/**/*.py[cod]
 pupy/*.py[cod]
-pupy/pupylib/*.py[cod]
+pupy/pupylib/**/*.py[cod]
 pupy/modules/*.py[cod]
 
 # do not ignore package & templates files

pupy/modules/memory_exec.py

@@ -15,6 +15,8 @@
 # --------------------------------------------------------------
 from pupylib.PupyModule import *
 from pupylib.PupyCompleter import *
+from pupylib.utils.pe import get_pe_arch
+from pupylib.PupyErrors import PupyModuleError
 
 __class_name__="MemoryExec"
 
@@ -30,7 +32,7 @@ class MemoryExec(PupyModule):
 		self.arg_parser.add_argument('-p', '--process', default='cmd.exe', help='process to start suspended')
 		self.arg_parser.add_argument('--fork', action='store_true', help='fork and do not wait for the child program. stdout will not be retrieved', completer=path_completer)
 		self.arg_parser.add_argument('--interactive', action='store_true', help='interactive with the new process stdin/stdout')
-		self.arg_parser.add_argument('path', help='path to the exe')
+		self.arg_parser.add_argument('path', help='path to the exe', completer=path_completer)
 		self.arg_parser.add_argument('args', nargs='*', help='optional arguments to pass to the exe')
 
 	@windows_only
@@ -49,6 +51,15 @@ class MemoryExec(PupyModule):
 			#TODO
 			self.error("interactive memory execution has not been implemented yet")
 			return
+		
+		#check we are injecting from the good process arch:
+		pe_arch=get_pe_arch(args.path)
+		proc_arch=self.client.desc["proc_arch"]
+		if pe_arch!=proc_arch:
+			self.error("%s is a %s PE and your pupy payload is a %s process. Please inject a %s PE or first migrate into a %s process"%(args.path, pe_arch, proc_arch, proc_arch, pe_arch))
+			return
+
+		
 		wait=True
 		redirect_stdio=True
 		if args.fork:

pupy/pupylib/utils/pe.py

@@ -0,0 +1,20 @@
+#!/usr/bin/env python
+# -*- coding: UTF8 -*-
+
+import pefile
+
+def get_pe_arch(*args, **kwargs):
+	pe=None
+	if args:
+		pe = pefile.PE(args[0], fast_load=True)
+	elif "data" in kwargs:
+		pe = pefile.PE(data=kwargs["data"], fast_load=True)
+	else:
+		raise NameError("at least a path or data must be supplied to get_arch")
+	if pe.OPTIONAL_HEADER.Magic==0x010b:
+		return "32bit"
+	elif pe.OPTIONAL_HEADER.Magic==0x020b:
+		return "64bit"
+	else:
+		return "UNKNOWN"
+