install_api.py: add argparse interface, refactoring, support pyenv.

1 files changed, 168 insertions, 101 deletions

diff --git a/scripts/install_api.py b/scripts/install_api.py
index 99f4db90..b571841a 100755
--- a/scripts/install_api.py
+++ b/scripts/install_api.py
@@ -1,133 +1,200 @@
 #!/usr/bin/env python3
-#
-#    Thanks to @withzombies for letting us adapt his script
-#
 
 import importlib.util
 import os
 import sys
 from site import check_enableusersite
+from typing import Optional
 
 if sys.version_info[0] < 3:
-	print("Python 3 required")
-	sys.exit(1)
+    sys.stderr.print("Python 3 is required.")
+    sys.exit(1)
 
 # Handle both normal environments and virtualenvs
 try:
-	from site import getusersitepackages, getsitepackages
+    from site import getusersitepackages, getsitepackages
 except ImportError:
-	from sysconfig import get_path
-	getsitepackages = lambda: get_path('purelib')
-	getusersitepackages = getsitepackages
+    from sysconfig import get_path
+    getsitepackages = lambda: get_path("purelib")
+    getusersitepackages = getsitepackages
 
-# Hacky command line parsing to accept a silent-install -s flag like linux-setup.sh:
-INTERACTIVE = True
-if '-s' in sys.argv[1:]:
-	INTERACTIVE = False
 
-# Autodetect venv
-INSTALL_VENV = False
-if sys.prefix != sys.base_prefix:
-	if os.environ.get('VIRTUAL_ENV'):
-		INSTALL_VENV = True
+def print_error(*args, file=None, **kwargs):
+    print(*args, file=sys.stderr, **kwargs)
 
-try:
-	binaryninja = importlib.util.find_spec('binaryninja')
-	assert binaryninja is not None
-	binaryninjaui = importlib.util.find_spec('binaryninjaui')
-	assert binaryninjaui is not None
 
-	if '-u' in sys.argv[1:]:
-		#Uninstall mode
-		userpth = os.path.join(getusersitepackages(), 'binaryninja.pth')
-		if os.path.exists(userpth):
-			print(f"Removing {userpth}")
-			os.unlink(userpth)
-		sitepth = os.path.join(getsitepackages()[0], 'binaryninja.pth')
-		if os.path.exists(sitepth):
-			print(f"Removing {sitepth}")
-			try:
-				os.unlink(sitepth)
-			except:
-				print("Unable to unlink, please re-run with appropriate permissions.")
-		sys.exit(0)
-	print("Binary Ninja API already in the path")
-	sys.exit(1)
-except ModuleNotFoundError:
-	pass
-except AssertionError:
-	pass
+def check_virtual_environment() -> bool:
+    if sys.prefix != sys.base_prefix:
+        if os.environ.get("VIRTUAL_ENV"):
+            return True
+    return False
+
+
+def binaryninja_installed() -> bool:
+    try:
+        binaryninja = importlib.util.find_spec("binaryninja")
+        assert binaryninja is not None
+        binaryninjaui = importlib.util.find_spec("binaryninjaui")
+        assert binaryninjaui is not None
+        return True
+    except ModuleNotFoundError:
+        return False
+    except AssertionError:
+        return False
+
+
+def get_binaryninja_installed_directory() -> Optional[str]:
+    dir_name = os.path.dirname(os.path.abspath(__file__))
+    api_path = os.path.abspath(os.path.join(dir_name, "..", "python"))
+    if os.path.isdir(api_path):
+        return api_path
+
+
+def validate_path(path: str) -> bool:
+
+    try:
+        os.stat(path)
+    except OSError:
+        return False
+
+    old_path = sys.path
+    sys.path.append(path)
+
+    try:
+        from binaryninja import core_version
+        print("Found Binary Ninja core version: {}".format(core_version()))
+    except ModuleNotFoundError as e:
+        sys.path = old_path
+        if e.name == "packaging" and sys.version_info.minor >= 10:
+            raise e
+        return False
+    except ImportError as e:
+        sys.path = old_path
+        return False
+
+    return True
+
+
+def install(interactive=False, on_root=False, on_pyenv=False) -> bool:
+
+    if binaryninja_installed():
+        print_error("Binary Ninja API already in the path.")
+        return False
+
+    api_path = get_binaryninja_installed_directory()
+    if not api_path:
+        print_error("Failed to find installed python expected at {}".format(api_path))
+        return False
+
+    print(f"Found install folder of {api_path}")
+
+    while not validate_path(api_path):
+        
+        print(f"Binary Ninja not found: {api_path}")
+        
+        if not interactive:
+            print_error("silent mode selected (-s, --silent), failing.")
+            return False
+
+        try:
+            new_path = input("Please provide the path to Binary Ninja's install directory: \n [{}] : ".format(api_path))
+        except KeyboardInterrupt:
+            print_error("KeyboardInterrupt detected.")
+            return False
+
+        if not new_path:
+            print("Invalid path.")
+            continue
+
+        if not new_path.endswith("python"):
+            os.path.join(new_path, "python")
+
+        api_path = new_path
+
+    if on_root:
+        install_path = getsitepackages()[0]
+        if not os.access(install_path, os.W_OK):
+            print_error(f"Root install specified but cannot write to {install_path}")
+            return False
+        else:
+            print(f"Installing on root site: {install_path}")
+    
+    elif on_pyenv:
+        install_path = getsitepackages()[0]
+        print(f"Installing on pyenv site: {install_path}")
 
-if '-u' in sys.argv[1:]:
-	print("Uninstall specified, but binaryninja not in the current path")
-	sys.exit(1)
+    elif check_virtual_environment():
+        install_path = getsitepackages()[0]
+        print(f"Installing on virtual environment site: {install_path}")
 
-dir_name = os.path.dirname(os.path.abspath(__file__))
-api_path = os.path.abspath(os.path.join(dir_name, "..", "python"))
+    else:
+        if not check_enableusersite():
+            print_error("Warning, trying to write to user site packages, but check_enableusersite fails.")
+            return False
+        else:
+            install_path = getusersitepackages()
+            if not os.path.exists(install_path):
+                os.makedirs(install_path)
+            print(f"Installing on user site: {install_path}")
 
-if (os.path.isdir(api_path)):
-	print(f"Found install folder of {api_path}")
-else:
-	print("Failed to find installed python expected at {}".format(api_path))
-	sys.exit(1)
+    binaryninja_pth_path = os.path.join(install_path, "binaryninja.pth")
+    with open(binaryninja_pth_path, 'wb') as pth_file:
+        pth_file.write((api_path + '\n').encode("charmap"))
+        pth_file.write((api_path + sys.version[0] + '\n').encode("charmap")) # support for QT bindings
 
+    return True
 
-def validate_path(path):
-	try:
-		os.stat(path)
-	except OSError:
-		return False
 
-	old_path = sys.path
-	sys.path.append(path)
+def uninstall() -> bool:
 
-	try:
-		from binaryninja import core_version
-		print("Found Binary Ninja core version: {}".format(core_version()))
-	except ImportError:
-		sys.path = old_path
-		return False
+    if not binaryninja_installed():
+        print_error("Uninstall specified, but binaryninja not in the current path.")
+        return False
 
-	return True
+    user_path = os.path.join(getusersitepackages(), "binaryninja.pth")
+    site_path = os.path.join(getsitepackages()[0], "binaryninja.pth")
 
+    for path_to_unlink in [user_path, site_path]:
 
-while not validate_path(api_path):
-	print("\nBinary Ninja not found.")
-	if not INTERACTIVE:
-		print("Non-interactive mode selected, failing.")
-		sys.exit(-1)
+        if not os.path.exists(path_to_unlink):
+            print_error(f"{path_to_unlink} not found.")
+            continue
 
-	print("Please provide the path to Binary Ninja's install directory: \n [{}] : ".format(api_path))
-	new_path = sys.stdin.readline().strip()
-	if len(new_path) == 0:
-		print("\nInvalid path")
-		continue
+        print(f"Removing {path_to_unlink}...")
+        try:
+            os.unlink(path_to_unlink)
+        except OSError as e:
+            print_error(f"Unable to unlink, please re-run with appropriate permissions: {str(e)}")
+            return False
 
-	if not new_path.endswith('python'):
-		new_path = os.path.join(new_path, 'python')
+    return True
 
-	api_path = new_path
 
-if (len(sys.argv) > 1 and sys.argv[1].lower() == "root"):
-	#write to root site
-	install_path = getsitepackages()[0]
-	if not os.access(install_path, os.W_OK):
-		print("Root install specified but cannot write to {}".format(install_path))
-		sys.exit(1)
-elif INSTALL_VENV:
-	install_path = getsitepackages()[0]
-else:
-	if check_enableusersite():
-		install_path = getusersitepackages()
-		if not os.path.exists(install_path):
-			os.makedirs(install_path)
-	else:
-		print("Warning, trying to write to user site packages, but check_enableusersite fails.")
-		sys.exit(1)
+if __name__ == '__main__':
 
-binaryninja_pth_path = os.path.join(install_path, 'binaryninja.pth')
-with open(binaryninja_pth_path, 'wb') as pth_file:
-	pth_file.write((api_path + "\n").encode('charmap'))
-	pth_file.write((api_path + sys.version[0] + "\n").encode('charmap'))  #support for QT bindings
+    import argparse
+    parser = argparse.ArgumentParser()
+    parser.add_argument("-s", "--silent", action='store_true') 
+    parser.add_argument("-u", "--uninstall", action='store_true')
+    parser.add_argument("--install-on-root", action='store_true')
+    parser.add_argument("--install-on-pyenv", action='store_true')
+    args = parser.parse_args()
 
-print("Binary Ninja API installed using {}".format(binaryninja_pth_path))
+    if args.uninstall:
+        uninstall_result = uninstall()
+        if not uninstall_result:
+            print(f"Binary Ninja API uninstallation failed.")
+            sys.exit(-1)
+        else:
+            print(f"Binary Ninja API uninstallation success.")
+    else:
+        install_result = install(
+            interactive=(not args.silent),
+            on_root=args.install_on_root,
+            on_pyenv=args.install_on_pyenv
+        )
+        if not install_result:
+            print(f"Binary Ninja API installation failed.")
+            sys.exit(-1)
+        else:
+            print(f"Binary Ninja API installation success.")