[Enterprise] Clean up initialization flow

5 files changed, 96 insertions, 51 deletions

diff --git a/python/__init__.py b/python/__init__.py
index 3c440c8b..ec218d38 100644
--- a/python/__init__.py
+++ b/python/__init__.py
@@ -216,24 +216,18 @@ class _DestructionCallbackHandler:
 
 _enable_default_log = True
 _plugin_init = False
+_enterprise_license_checkout = None
 
 
 def _init_plugins():
 	global _enable_default_log
 	global _plugin_init
+	global _enterprise_license_checkout
 
 	if not core_ui_enabled() and core.BNGetProduct() == "Binary Ninja Enterprise Client":
 		# Enterprise client needs to checkout a license reservation or else BNInitPlugins will fail
-		if not enterprise.is_license_still_activated():
-			try:
-				enterprise.authenticate_with_method("Keychain")
-			except RuntimeError:
-				pass
-		if not core.BNIsLicenseValidated() or not enterprise.is_license_still_activated():
-			raise RuntimeError(
-			    "To use Binary Ninja Enterprise from a headless python script, you must check out a license first.\n"
-			    "You can either check out a license for an extended time with the UI, or use the binaryninja.enterprise module."
-			)
+		_enterprise_license_checkout = enterprise.LicenseCheckout()
+		_enterprise_license_checkout.acquire()
 
 	if not _plugin_init:
 		# The first call to BNInitCorePlugins returns True for successful initialization and True in this context indicates headless operation.
diff --git a/python/collaboration/__init__.py b/python/collaboration/__init__.py
index 63331d83..dd20eebb 100644
--- a/python/collaboration/__init__.py
+++ b/python/collaboration/__init__.py
@@ -75,14 +75,15 @@ def enterprise_remote() -> Optional['Remote']:
 	return None
 
 
-def add_known_remote(remote: 'Remote') -> None:
+def create_remote(name: str, address: str) -> 'Remote':
 	"""
-	Add a Remote to the list of known remotes (saved to Settings)
+	Create a Remote and add it to the list of known remotes (saved to Settings)
 
-	:param remote: New Remote to add
+	:param name: Identifier for remote
+	:param address: Base address (HTTPS) for all api requests
 	"""
 	binaryninja._init_plugins()
-	core.BNCollaborationAddRemote(remote._handle)
+	return Remote(core.BNCollaborationCreateRemote(name, address))
 
 
 def remove_known_remote(remote: 'Remote') -> None:
diff --git a/python/collaboration/remote.py b/python/collaboration/remote.py
index dbd72598..811fd33b 100644
--- a/python/collaboration/remote.py
+++ b/python/collaboration/remote.py
@@ -1,5 +1,6 @@
 import ctypes
 import json
+import os
 from typing import Dict, List, Optional, Tuple
 
 import binaryninja
@@ -15,10 +16,6 @@ class Remote:
 	"""
 	def __init__(self, handle: core.BNRemoteHandle):
 		"""
-		Create a Remote object (but don't connect to it yet)
-
-		:param name: Identifier for remote
-		:param address: Base address (HTTPS) for all api requests
 		:param handle: FFI handle for internal use
 		:raises: RuntimeError if there was an error
 		"""
@@ -245,22 +242,37 @@ class Remote:
 		"""
 		if not self.has_loaded_metadata:
 			self.load_metadata()
-		if username is None:
+		got_auth = False
+		if username is not None and token is not None:
+			got_auth = True
+		if not got_auth:
 			# Try logging in with defaults
-			if self.is_enterprise:
+			if self.is_enterprise and enterprise.is_authenticated():
 				username = enterprise.username()
 				token = enterprise.token()
-			else:
-				# Load from default secrets provider
-				secrets = binaryninja.SecretsProvider[
-					binaryninja.Settings().get_string("enterprise.secretsProvider")]
-				if not secrets.has_data(self.address):
-					raise RuntimeError("No username and token provided, and none found "
-									   "in the default keychain.")
+				if username is not None and token is not None:
+					got_auth = True
+
+		if not got_auth:
+			# Try to load from default secrets provider
+			secrets = binaryninja.SecretsProvider[
+				binaryninja.Settings().get_string("enterprise.secretsProvider")]
+			if secrets.has_data(self.address):
 				creds = json.loads(secrets.get_data(self.address))
 				username = creds['username']
 				token = creds['token']
-		if username is None or token is None:
+				got_auth = True
+
+		if not got_auth:
+			# Try logging in with creds in the env
+			if os.environ.get('BN_ENTERPRISE_USERNAME') is not None and \
+				os.environ.get('BN_ENTERPRISE_PASSWORD') is not None:
+				token = self.request_authentication_token(os.environ['BN_ENTERPRISE_USERNAME'], os.environ['BN_ENTERPRISE_PASSWORD'])
+				if token is not None:
+					username = os.environ['BN_ENTERPRISE_USERNAME']
+					got_auth = True
+
+		if not got_auth or username is None or token is None:
 			raise RuntimeError("Cannot connect without a username or token")
 
 		if not core.BNRemoteConnect(self._handle, username, token):
diff --git a/python/enterprise.py b/python/enterprise.py
index e645ed5f..5b4403d5 100644
--- a/python/enterprise.py
+++ b/python/enterprise.py
@@ -18,6 +18,23 @@ if core.BNGetProduct() != "Binary Ninja Enterprise Client":
 	raise RuntimeError("Cannot use Binary Ninja Enterprise client functionality with a non-Enterprise client.")
 
 
+def is_initialized() -> bool:
+	"""
+	Determine if the Enterprise Client has been initialized yet.
+
+	:return: True if :py:func:`initialize` has been called
+	"""
+	return core.BNIsEnterpriseServerInitialized()
+
+
+def initialize():
+	"""
+	Initialize the Enterprise Client
+	"""
+	if not core.BNInitializeEnterpriseServer():
+		raise RuntimeError(last_error())
+
+
 def connect():
 	"""
 	Connect to the Enterprise Server.
@@ -117,8 +134,8 @@ def username() -> Optional[str]:
 	:return: Username, if authenticated. None, otherwise.
 	"""
 	value = core.BNGetEnterpriseServerUsername()
-	if value == "":
-		return None
+	if value is None:
+		raise RuntimeError(last_error())
 	return value
 
 
@@ -129,8 +146,8 @@ def token() -> Optional[str]:
 	:return: Token, if authenticated. None, otherwise.
 	"""
 	value = core.BNGetEnterpriseServerToken()
-	if value == "":
-		return None
+	if value is None:
+		raise RuntimeError(last_error())
 	return value
 
 
@@ -140,7 +157,10 @@ def server_url() -> str:
 
 	:return: The current url
 	"""
-	return core.BNGetEnterpriseServerUrl()
+	value = core.BNGetEnterpriseServerUrl()
+	if value is None:
+		raise RuntimeError(last_error())
+	return value
 
 
 def set_server_url(url: str):
@@ -164,8 +184,8 @@ def server_name() -> str:
 	if not is_connected():
 		connect()
 	value = core.BNGetEnterpriseServerName()
-	if value == "":
-		return None
+	if value is None:
+		raise RuntimeError(last_error())
 	return value
 
 
@@ -178,8 +198,8 @@ def server_id() -> str:
 	if not is_connected():
 		connect()
 	value = core.BNGetEnterpriseServerId()
-	if value == "":
-		return None
+	if value is None:
+		raise RuntimeError(last_error())
 	return value
 
 
@@ -193,7 +213,7 @@ def server_version() -> int:
 		connect()
 	value = core.BNGetEnterpriseServerVersion()
 	if value == 0:
-		return None
+		raise RuntimeError(last_error())
 	return value
 
 
@@ -206,8 +226,8 @@ def server_build_id() -> str:
 	if not is_connected():
 		connect()
 	value = core.BNGetEnterpriseServerBuildId()
-	if value == "":
-		return None
+	if value is None:
+		raise RuntimeError(last_error())
 	return value
 
 
@@ -234,6 +254,7 @@ def update_license(duration, _cache=True):
 	if not core.BNUpdateEnterpriseServerLicense(duration):
 		raise RuntimeError(last_error())
 
+
 @deprecation.deprecated(deprecated_in="3.4.4137", details="Use .update_license instead.")
 def acquire_license(duration, _cache=True):
 	"""
@@ -248,6 +269,7 @@ def acquire_license(duration, _cache=True):
 	"""
 	update_license(duration, _cache)
 
+
 def release_license():
 	"""
 	Release the currently checked out license back to the Enterprise Server.
@@ -306,15 +328,6 @@ def last_error() -> str:
 	return core.BNGetEnterpriseServerLastError()
 
 
-def is_initialized() -> bool:
-	"""
-	Determine if the Enterprise Client has been initialized yet.
-
-	:return: True if any other Enterprise methods have been called
-	"""
-	return core.BNIsEnterpriseServerInitialized()
-
-
 @decorators.enterprise
 class LicenseCheckout:
 	"""
@@ -347,10 +360,31 @@ class LicenseCheckout:
 		self.acquired_license = False
 		self.desired_release = release
 
+	def __del__(self):
+		self.release()
+
 	def __enter__(self) -> None:
+		self.acquire()
+
+	def __exit__(self, exc_type, exc_val, exc_tb):
+		self.release()
+
+	def acquire(self):
 		# UI builds have their own license manager
 		if binaryninja.core_ui_enabled():
 			return
+		if not is_initialized():
+			try:
+				initialize()
+			except:
+				# Named/computer licenses don't need this flow at all
+				if not is_floating_license():
+					return
+				# Floating licenses though, this is an error. Probably the error
+				# for needing to set enterprise.server.url in settings.json
+				raise
+		if not is_floating_license():
+			return
 		if not is_connected():
 			connect()
 		got_auth = False
@@ -385,10 +419,11 @@ class LicenseCheckout:
 			update_license(self.desired_duration)
 			self.acquired_license = True
 
-	def __exit__(self, exc_type, exc_val, exc_tb):
+	def release(self):
 		# UI builds have their own license manager
 		if binaryninja.core_ui_enabled():
 			return
 		# Don't release if we got one from keychain
 		if self.acquired_license and self.desired_release:
 			release_license()
+			self.acquired_license = False
diff --git a/python/generator.cpp b/python/generator.cpp
index 3fdf1fc1..c40d8d08 100644
--- a/python/generator.cpp
+++ b/python/generator.cpp
@@ -567,7 +567,10 @@ int main(int argc, char* argv[])
 			// Emit wrapper to get Python string and free native memory
 			fprintf(out, "\tresult = ");
 			fprintf(out, "%s\n", stringArgFuncCall.c_str());
-			fprintf(out, "\tstring = str(pyNativeStr(ctypes.cast(result, ctypes.c_char_p).value))\n");
+			fprintf(out, "\tcasted = ctypes.cast(result, ctypes.c_char_p).value\n");
+			fprintf(out, "\tif casted is None:\n");
+			fprintf(out, "\t\treturn None\n");
+			fprintf(out, "\tstring = str(pyNativeStr(casted))\n");
 			fprintf(out, "\tBNFreeString(result)\n");
 			fprintf(out, "\treturn string\n");
 		}