Enterprise: Add AuthenticateWithToken api

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

Author: CouleeApps

binaryninjacore.h

@@ -4063,6 +4063,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);

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);
@@ -281,11 +287,20 @@ 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"
 			);
 		}

enterprise.h

@@ -52,6 +52,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

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 @@ def acquire(self):
 				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"
 				)
 

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,