Updated setup.py

setup.py

@@ -2,41 +2,45 @@
 from __future__ import print_function
 import os
 import sys
-from distutils.core import setup
+import tempfile
+import subprocess
+from distutils.core import setup, Command
 from distutils.dir_util import remove_tree
 
+MODULE_NAME = "binwalk"
+
 # Python2/3 compliance
 try:
 	raw_input
 except NameError:
 	raw_input = input
 
-def cleanup_build_directory():
-	# Requires to chdir into the src directory first
-	try:
-		remove_tree("build")
-	except KeyboardInterrupt as e:
-		raise e
-	except Exception:
-		pass
+# This is super hacky.
+if "--yes" in sys.argv:
+	sys.argv.pop(sys.argv.index("--yes"))
+	IGNORE_WARNINGS = True
+else:
+	IGNORE_WARNINGS = False
+
+# cd into the src directory, no matter where setup.py was invoked from
+os.chdir(os.path.join(os.path.dirname(os.path.realpath(__file__)), "src"))
+
+def which(command):
+	# /usr/local/bin is usually the default install path, though it may not be in $PATH
+	usr_local_bin = os.path.sep.join([os.path.sep, 'usr', 'local', 'bin', command])
 
-def cleanup_module_directory():
-	# Installing doesn't remove old files that may have been deleted from the module.
-	if "install" in sys.argv:
-		try:
-			import binwalk
-			for path in binwalk.__path__:
 	try:
-					remove_tree(path + os.path.sep + "*")
-				except OSError as e:
-					pass
+		location = subprocess.Popen(["which", command], shell=False, stdout=subprocess.PIPE).communicate()[0].strip()
 	except KeyboardInterrupt as e:
 		raise e
-		except Exception:
+	except Exception as e:
 		pass
 
+	if not location and os.path.exists(usr_local_bin):
+		location = usr_local_bin
+
+	return location
 
-# Change to the binwalk src directory
 def warning(lines, terminate=True, prompt=True):
 	WIDTH = 115
 
@@ -55,102 +59,105 @@ def warning(lines, terminate=True, prompt=True):
 		if terminate:
 			sys.exit(1)
 
-# This is super hacky.
-if "--yes" in sys.argv:
-	sys.argv.pop(sys.argv.index("--yes"))
-	IGNORE_WARNINGS = True
-else:
-	IGNORE_WARNINGS = False
-
-# cd into the src directory, no matter where setup.py was invoked from
-os.chdir(os.path.join(os.path.dirname(os.path.realpath(__file__)), "src"))
+class UninstallCommand(Command):
+	description = "Uninstalls the Python module"
+	user_options = []
 
-print("checking pre-requisites")
-try:
-	import magic
-	try:
-		magic.MAGIC_NO_CHECK_TEXT
-	except AttributeError as e:
-		msg = ["Pre-requisite failure: " + str(e),
-			"It looks like you have an old or incompatible magic module installed.",
-			"Please install the official python-magic module, or download and install it from source: ftp://ftp.astron.com/pub/file/"
-		]
+	def initialize_options(self):
+		pass
 
-		warning(msg)
-except ImportError as e:
-	msg = ["Pre-requisite failure:", str(e),
-		"Please install the python-magic module, or download and install it from source: ftp://ftp.astron.com/pub/file/",
-	]
+	def finalize_options(self):
+		pass
 
-	warning(msg)
+	def run(self):
+		try:
+			import binwalk
+			for path in binwalk.__path__:
+				try:
+					remove_tree(path)
+				except OSError as e:
+					pass
+		except KeyboardInterrupt as e:
+			raise e
+		except Exception as e:
+			pass
 
-try:
-	import pyqtgraph
-except ImportError as e:
-	msg = ["Pre-requisite check warning: " + str(e),
-		"To take advantage of this tool's graphing capabilities, please install the pyqtgraph module.",
-	]
+		script_path = which(MODULE_NAME)
+		if script_path:
+			try:
+				print("removing '%s'" % script_path)
+				os.unlink(script_path)
+			except KeyboardInterrupt as e:
+				pass
+			except Exception as e:
+				pass
 
-	warning(msg, prompt=True)
+class CleanCommand(Command):
+	description = "Clean Python build directories"
+	user_options = []
 
-# Build / install C compression libraries
-c_lib_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), "C")
-c_lib_makefile = os.path.join(c_lib_dir, "Makefile")
+	def initialize_options(self):
+		pass
 
-working_directory = os.getcwd()
-os.chdir(c_lib_dir)
-status = 0
+	def finalize_options(self):
+		pass
 
-if not os.path.exists(c_lib_makefile):
-	status |= os.system("./configure")
+	def run(self):
+		try:
+			remove_tree("build")
+		except KeyboardInterrupt as e:
+			raise e
+		except Exception:
+			pass
 
-status |= os.system("make")
+		try:
+			remove_tree("dist")
+		except KeyboardInterrupt as e:
+			raise e
+		except Exception:
+			pass
 
-if status != 0:
-	msg = ["Build warning: failed to build C libraries.",
-		"Some features will not work without these libraries."
-	]
-	warning(msg, prompt=True)
-elif "install" in sys.argv:
-	if os.system("make install") != 0:
-		msg = ["Install warning: failed to install C libraries.",
-			   "Some features will not work without these libraries."
+# Check pre-requisite Python modules during an install
+if "install" in sys.argv:
+	print("checking pre-requisites")
+	try:
+		import pyqtgraph
+		import QtPyWhatever
+	except ImportError as e:
+		msg = ["Pre-requisite check warning: " + str(e),
+			   "To take advantage of this %s's graphing capabilities, please install this module." % MODULE_NAME,
 		]
-		warning(msg, prompt=True)
 	
-os.system("make distclean")
-	
-os.chdir(working_directory)
-
-cleanup_build_directory()
-cleanup_module_directory()
+		warning(msg, prompt=True)
 
-# Generate a new magic file from the files in the magic directory
-print("creating binwalk magic file")
-magic_files = os.listdir("magic")
-magic_files.sort()
-fd = open("binwalk/magic/binwalk", "wb")
-for magic in magic_files:
+# Re-build the magic file during a build/install
+if "install" in sys.argv or "build" in sys.argv:
+	# Generate a new magic file from the files in the magic directory
+	print("creating %s magic file" % MODULE_NAME)
+	magic_files = os.listdir("magic")
+	magic_files.sort()
+	fd = open("%s/magic/%s" % (MODULE_NAME, MODULE_NAME), "wb")
+	for magic in magic_files:
 		fpath = os.path.join("magic", magic)
 		if os.path.isfile(fpath):
 			fd.write(open(fpath, "rb").read())
-fd.close()
+	fd.close()
 
-# The data files to install along with the binwalk module
+# The data files to install along with the module
 install_data_files = ["magic/*", "config/*", "plugins/*", "modules/*", "core/*"]
 
-# Install the binwalk module, script and support files
-setup(	name = "binwalk",
+# Install the module, script, and support files
+setup(name = MODULE_NAME,
 	  version = "2.0.0 beta",
 	  description = "Firmware analysis tool",
 	  author = "Craig Heffner",
-	url = "https://github.com/devttys0/binwalk",
+	  url = "https://github.com/devttys0/%s" % MODULE_NAME,
 
 	  requires = ["magic", "pyqtgraph"],
-	packages = ["binwalk"],
-	package_data = {"binwalk" : install_data_files},
-	scripts = ["scripts/binwalk"],
-)
+	  packages = [MODULE_NAME],
+	  package_data = {MODULE_NAME : install_data_files},
+	  scripts = [os.path.join("scripts", MODULE_NAME)],
 
-cleanup_build_directory()
+	  cmdclass = {'clean' : CleanCommand, 'uninstall' : UninstallCommand}
+)