[Enterprise] Clean up initialization flow

Author: CouleeApps

binaryninjacore.h

@@ -3359,6 +3359,7 @@ extern "C"
 	BINARYNINJACOREAPI void BNRegisterEnterpriseServerNotification(BNEnterpriseServerCallbacks* notify);
 	BINARYNINJACOREAPI void BNUnregisterEnterpriseServerNotification(BNEnterpriseServerCallbacks* notify);
 	BINARYNINJACOREAPI bool BNIsEnterpriseServerInitialized(void);
+	BINARYNINJACOREAPI bool BNInitializeEnterpriseServer(void);
 
 	BINARYNINJACOREAPI void BNRegisterObjectDestructionCallbacks(BNObjectDestructionCallbacks* callbacks);
 	BINARYNINJACOREAPI void BNUnregisterObjectDestructionCallbacks(BNObjectDestructionCallbacks* callbacks);

enterprise.cpp

@@ -23,6 +23,19 @@
 
 using namespace BinaryNinja::Enterprise;
 
+
+bool BinaryNinja::Enterprise::IsInitialized()
+{
+	return BNIsEnterpriseServerInitialized();
+}
+
+
+bool BinaryNinja::Enterprise::Initialize()
+{
+	return BNInitializeEnterpriseServer();
+}
+
+
 bool BinaryNinja::Enterprise::AuthenticateWithCredentials(const std::string& username, const std::string& password, bool remember)
 {
 	return BNAuthenticateEnterpriseServerWithCredentials(username.c_str(), password.c_str(), remember);
@@ -99,6 +112,8 @@ bool BinaryNinja::Enterprise::IsAuthenticated()
 std::string BinaryNinja::Enterprise::GetUsername()
 {
 	char* value = BNGetEnterpriseServerUsername();
+	if (!value)
+		return "";
 	std::string result = value;
 	BNFreeString(value);
 	return result;
@@ -108,6 +123,8 @@ std::string BinaryNinja::Enterprise::GetUsername()
 std::string BinaryNinja::Enterprise::GetToken()
 {
 	char* value = BNGetEnterpriseServerToken();
+	if (!value)
+		return "";
 	std::string result = value;
 	BNFreeString(value);
 	return result;
@@ -117,6 +134,8 @@ std::string BinaryNinja::Enterprise::GetToken()
 std::string BinaryNinja::Enterprise::GetServerName()
 {
 	char* value = BNGetEnterpriseServerName();
+	if (!value)
+		return "";
 	std::string result = value;
 	BNFreeString(value);
 	return result;
@@ -126,6 +145,8 @@ std::string BinaryNinja::Enterprise::GetServerName()
 std::string BinaryNinja::Enterprise::GetServerUrl()
 {
 	char* value = BNGetEnterpriseServerUrl();
+	if (!value)
+		return "";
 	std::string result = value;
 	BNFreeString(value);
 	return result;
@@ -135,6 +156,8 @@ std::string BinaryNinja::Enterprise::GetServerUrl()
 std::string BinaryNinja::Enterprise::GetServerId()
 {
 	char* value = BNGetEnterpriseServerId();
+	if (!value)
+		return "";
 	std::string result = value;
 	BNFreeString(value);
 	return result;
@@ -150,6 +173,8 @@ uint64_t BinaryNinja::Enterprise::GetServerVersion()
 std::string BinaryNinja::Enterprise::GetServerBuildId()
 {
 	char* value = BNGetEnterpriseServerBuildId();
+	if (!value)
+		return "";
 	std::string result = value;
 	BNFreeString(value);
 	return result;
@@ -188,7 +213,6 @@ bool BinaryNinja::Enterprise::IsLicenseStillActivated()
 
 std::string BinaryNinja::Enterprise::GetLastError()
 {
-	return BNGetEnterpriseServerLastError();
 	char* str = BNGetEnterpriseServerLastError();
 	std::string value = str;
 	BNFreeString(str);
@@ -208,14 +232,36 @@ void BinaryNinja::Enterprise::UnregisterNotification(BNEnterpriseServerCallbacks
 }
 
 
-BinaryNinja::Enterprise::LicenseCheckout::LicenseCheckout(int64_t duration)
+BinaryNinja::Enterprise::LicenseCheckout::LicenseCheckout(int64_t duration): m_acquiredLicense(false)
 {
 	// This is a port of python's binaryninja.enterprise.LicenseCheckout
 
 	// UI builds have their own license manager
 	if (BNIsUIEnabled())
 		return;
 
+	if (!IsInitialized())
+	{
+		if (!Initialize())
+		{
+			// Named/computer licenses don't need this flow at all
+			if (!IsFloatingLicense())
+			{
+				return;
+			}
+
+			// Floating licenses though, this is an error. Probably the error
+			// for needing to set enterprise.server.url in settings.json
+			throw EnterpriseException(
+				"Could not initialize Binary Ninja Enterprise: " + GetLastError()
+			);
+		}
+	}
+
+	if (!IsFloatingLicense())
+	{
+		return;
+	}
 	if (!IsConnected())
 	{
 		Connect();

enterprise.h

@@ -40,6 +40,18 @@ namespace BinaryNinja
 			EnterpriseException(const std::string& what): std::runtime_error(what) {}
 		};
 
+		/*!
+		    Determine if the Enterprise Client has been initialized yet.
+		    \return True if Initialize() has been called successfully
+		 */
+		bool IsInitialized();
+
+		/*!
+		    Initialize the Enterprise Client
+		    \return True if successful
+		 */
+		bool Initialize();
+
 		/*!
 		    Authenticate to the Enterprise server with username and password
 		    \param username Username to authenticate with

python/__init__.py

@@ -216,24 +216,18 @@ def destruct_function(self, ctxt, func):
 
 _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.

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:

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 @@ def connect(self, username: Optional[str] = None, token: Optional[str] = None):
 		"""
 		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):