Enterprise: Add AuthenticateWithToken api

For when you have a token and want to auth with it directly

5 files changed, 87 insertions, 12 deletions

diff --git a/binaryninjacore.h b/binaryninjacore.h
index a0b985ec..879b7e1e 100644
--- a/binaryninjacore.h
+++ b/binaryninjacore.h
@@ -4065,6 +4065,8 @@ extern "C"
 	BINARYNINJACOREAPI bool BNIsDatabase(const char* filename);
 	BINARYNINJACOREAPI bool BNIsDatabaseFromData(const void* data, size_t len);
 
+	BINARYNINJACOREAPI bool BNAuthenticateEnterpriseServerWithToken(
+	    const char* token, bool remember);
 	BINARYNINJACOREAPI bool BNAuthenticateEnterpriseServerWithCredentials(
 	    const char* username, const char* password, bool remember);
 	BINARYNINJACOREAPI bool BNAuthenticateEnterpriseServerWithMethod(const char* method, bool remember);
diff --git a/enterprise.cpp b/enterprise.cpp
index d2f9ccad..ba3b5db7 100644
--- a/enterprise.cpp
+++ b/enterprise.cpp
@@ -36,6 +36,12 @@ bool BinaryNinja::Enterprise::Initialize()
 }
 
 
+bool BinaryNinja::Enterprise::AuthenticateWithToken(const std::string& token, bool remember)
+{
+	return BNAuthenticateEnterpriseServerWithToken(token.c_str(), remember);
+}
+
+
 bool BinaryNinja::Enterprise::AuthenticateWithCredentials(const std::string& username, const std::string& password, bool remember)
 {
 	return BNAuthenticateEnterpriseServerWithCredentials(username.c_str(), password.c_str(), remember);
@@ -282,10 +288,19 @@ BinaryNinja::Enterprise::LicenseCheckout::LicenseCheckout(int64_t duration): m_a
 		}
 		if (!gotAuth)
 		{
+			char* token = getenv("BN_ENTERPRISE_TOKEN");
+			if (token)
+			{
+				gotAuth = AuthenticateWithToken(token, true);
+			}
+		}
+		if (!gotAuth)
+		{
 			throw EnterpriseException(
 				"Could not checkout a license: Not authenticated. Try one of the following: \n"
 				" - Log in and check out a license for an extended time\n"
 				" - Set BN_ENTERPRISE_USERNAME and BN_ENTERPRISE_PASSWORD environment variables\n"
+				" - Set BN_ENTERPRISE_TOKEN environment variable\n"
 				" - Use BinaryNinja::Enterprise::AuthenticateWithCredentials or AuthenticateWithMethod in your code"
 			);
 		}
diff --git a/enterprise.h b/enterprise.h
index cf76edb0..e4a5695e 100644
--- a/enterprise.h
+++ b/enterprise.h
@@ -53,6 +53,14 @@ namespace BinaryNinja
 		bool Initialize();
 
 		/*!
+		    Authenticate to the server with an access token
+		    \param token Token of auth session
+		    \param remember Remember token in keychain
+		    \return True if successful
+		 */
+		bool AuthenticateWithToken(const std::string& token, bool remember);
+
+		/*!
 		    Authenticate to the Enterprise server with username and password
 		    \param username Username to authenticate with
 		    \param password Password to authenticate with
diff --git a/python/enterprise.py b/python/enterprise.py
index 4bc21d9b..b11be1dd 100644
--- a/python/enterprise.py
+++ b/python/enterprise.py
@@ -54,6 +54,19 @@ def is_connected() -> bool:
 	return core.BNIsEnterpriseServerConnected()
 
 
+def authenticate_with_token(token: str, remember: bool = True):
+	"""
+	Authenticate to the server with an access token
+
+	:param str token: Token of auth session.
+	:param bool remember: Remember token in keychain
+	"""
+	if not is_connected():
+		connect()
+	if not core.BNAuthenticateEnterpriseServerWithToken(token, remember):
+		raise RuntimeError(last_error())
+
+
 def authenticate_with_credentials(username: str, password: str, remember: bool = True):
 	"""
 	Authenticate to the Enterprise Server with username/password credentials.
@@ -392,11 +405,20 @@ class LicenseCheckout:
 				except RuntimeError:
 					pass
 
+			if not got_auth and \
+				os.environ.get('BN_ENTERPRISE_TOKEN') is not None:
+				try:
+					authenticate_with_token(os.environ['BN_ENTERPRISE_TOKEN'])
+					got_auth = True
+				except RuntimeError:
+					pass
+
 			if not got_auth:
 				raise RuntimeError(
 					"Could not checkout a license: Not authenticated. Try one of the following: \n"
 					" - Log in and check out a license for an extended time\n"
 					" - Set BN_ENTERPRISE_USERNAME and BN_ENTERPRISE_PASSWORD environment variables\n"
+					" - Set BN_ENTERPRISE_TOKEN environment variable\n"
 					" - Use binaryninja.enterprise.authenticate_with_credentials or authenticate_with_method in your code"
 				)
 
diff --git a/rust/src/enterprise.rs b/rust/src/enterprise.rs
index 68ad069c..2e132454 100644
--- a/rust/src/enterprise.rs
+++ b/rust/src/enterprise.rs
@@ -9,10 +9,8 @@ use thiserror::Error;
 pub enum EnterpriseCheckoutError {
     #[error("enterprise server returned error: {0}")]
     ServerError(String),
-    #[error("no username set for credential authentication")]
-    NoUsername,
-    #[error("no password set for credential authentication")]
-    NoPassword,
+    #[error("no credentials set for authentication")]
+    NoCredentials,
     #[error("failed to authenticate with username and password")]
     NotAuthenticated,
     #[error("failed to refresh expired license: {0}")]
@@ -59,16 +57,36 @@ pub fn checkout_license(
 
         #[allow(clippy::collapsible_if)]
         if !is_server_authenticated() {
-            // We have yet to authenticate with the server, we should try all available authentication methods.
-            if !authenticate_server_with_method("Keychain", false) {
+            'auth: {
+                // We have yet to authenticate with the server, we should try all available authentication methods.
+                if authenticate_server_with_method("Keychain", false) {
+                    break 'auth;
+                }
+
                 // We could not authenticate with the system keychain, we should try with credentials.
-                let username = std::env::var("BN_ENTERPRISE_USERNAME")
-                    .map_err(|_| EnterpriseCheckoutError::NoUsername)?;
-                let password = std::env::var("BN_ENTERPRISE_PASSWORD")
-                    .map_err(|_| EnterpriseCheckoutError::NoPassword)?;
-                if !authenticate_server_with_credentials(&username, &password, true) {
-                    return Err(EnterpriseCheckoutError::NotAuthenticated);
+                let username = std::env::var("BN_ENTERPRISE_USERNAME");
+                let password = std::env::var("BN_ENTERPRISE_PASSWORD");
+                if let Ok(username) = username {
+                    if let Ok(password) = password {
+                        // Having creds that don't work is a hard error
+                        if !authenticate_server_with_credentials(&username, &password, true) {
+                            return Err(EnterpriseCheckoutError::NotAuthenticated);
+                        }
+                        // Otherwise, if the creds worked, we got auth
+                        break 'auth;
+                    }
+                }
+
+                let token = std::env::var("BN_ENTERPRISE_TOKEN");
+                if let Ok(token) = token {
+                    if !authenticate_server_with_token(&token, true) {
+                        return Err(EnterpriseCheckoutError::NotAuthenticated);
+                    }
+                    break 'auth;
                 }
+
+                // If we're still here, we don't have any credentials for authentication.
+                return Err(EnterpriseCheckoutError::NoCredentials);
             }
         }
     }
@@ -183,6 +201,16 @@ pub fn is_server_license_still_activated() -> bool {
     unsafe { binaryninjacore_sys::BNIsEnterpriseServerLicenseStillActivated() }
 }
 
+pub fn authenticate_server_with_token(token: &str, remember: bool) -> bool {
+    let token = token.to_cstr();
+    unsafe {
+        binaryninjacore_sys::BNAuthenticateEnterpriseServerWithToken(
+            token.as_ref().as_ptr() as *const std::os::raw::c_char,
+            remember,
+        )
+    }
+}
+
 pub fn authenticate_server_with_credentials(
     username: &str,
     password: &str,