From 855dd6352f3003ecd3062243769ef5bf5c303e54 Mon Sep 17 00:00:00 2001
From: Mason Reed <mason@vector35.com>
Date: Wed, 2 Oct 2024 19:51:27 -0400
Subject: [PATCH] Checkout enterprise license in rust headless mode

---
 rust/src/enterprise.rs | 66 ++++++++++++++++++++++++++++++++++++++++++
 rust/src/headless.rs   | 15 ++++++++++
 2 files changed, 81 insertions(+)

diff --git a/rust/src/enterprise.rs b/rust/src/enterprise.rs
index 232526dcb..5b8877f3c 100644
--- a/rust/src/enterprise.rs
+++ b/rust/src/enterprise.rs
@@ -4,6 +4,68 @@ use std::time::{Duration, SystemTime, UNIX_EPOCH};
 use crate::rc::Array;
 use crate::string::{BnStrCompatible, BnString};
 
+
+#[derive(Debug)]
+pub struct EnterpriseCheckoutError(pub String);
+
+impl std::fmt::Display for EnterpriseCheckoutError {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{}", self.0)
+    }
+}
+
+impl std::error::Error for EnterpriseCheckoutError {}
+
+pub fn checkout_license(duration: Duration) -> Result<(), EnterpriseCheckoutError> {
+    if crate::is_ui_enabled() {
+        return Ok(());
+    }
+
+    if !is_server_initialized() {
+        if !initialize_server() && is_server_floating_license() {
+            return Err(EnterpriseCheckoutError(server_last_error().to_string()));
+        }
+    }
+
+    if is_server_floating_license() {
+        if !is_server_connected() && !connect_server() {
+            return Err(EnterpriseCheckoutError(server_last_error().to_string()));
+        }
+
+        if !is_server_authenticated() {
+            if !authenticate_server_with_method("Keychain", false) {
+                let Some(username) = std::env::var("BN_ENTERPRISE_USERNAME").ok() else {
+                    return Err(EnterpriseCheckoutError("BN_ENTERPRISE_USERNAME not set when attempting to authenticate with credentials".to_string()));
+                };
+                let Some(password) = std::env::var("BN_ENTERPRISE_PASSWORD").ok() else {
+                    return Err(EnterpriseCheckoutError("BN_ENTERPRISE_PASSWORD not set when attempting to authenticate with credentials".to_string()));
+                };
+                if !authenticate_server_with_credentials(username, password, true) {
+                    let failed_message = "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 \
+                         - Use binaryninja::enterprise::{authenticate_server_with_method OR authenticate_server_with_credentials} in your code";
+                    return Err(EnterpriseCheckoutError(failed_message.to_string()));
+                }
+            }
+        }
+    }
+    
+    if !is_server_license_still_activated() || (!is_server_floating_license() && crate::license_expiration_time() < SystemTime::now()) {
+        if !update_server_license(duration) {
+            return Err(EnterpriseCheckoutError("Failed to refresh expired license".to_string()));
+        }
+    }
+
+    Ok(())
+}
+
+pub fn release_license() {
+    if !crate::is_ui_enabled() {
+        release_server_license();
+    }
+}
+
 pub fn server_username() -> BnString {
     unsafe { BnString::from_raw(binaryninjacore_sys::BNGetEnterpriseServerUsername()) }
 }
@@ -125,6 +187,10 @@ pub fn is_server_initialized() -> bool {
     unsafe { binaryninjacore_sys::BNIsEnterpriseServerInitialized() }
 }
 
+pub fn initialize_server() -> bool {
+    unsafe { binaryninjacore_sys::BNInitializeEnterpriseServer() }
+}
+
 pub fn server_last_error() -> BnString {
     unsafe { BnString::from_raw(binaryninjacore_sys::BNGetEnterpriseServerLastError()) }
 }
diff --git a/rust/src/headless.rs b/rust/src/headless.rs
index 1de41aecc..0cc92faaf 100644
--- a/rust/src/headless.rs
+++ b/rust/src/headless.rs
@@ -20,6 +20,7 @@ use crate::{
 
 use std::env;
 use std::path::PathBuf;
+use std::time::Duration;
 
 #[cfg(not(target_os = "windows"))]
 fn binja_path() -> PathBuf {
@@ -82,6 +83,13 @@ use binaryninjacore_sys::{BNInitPlugins, BNInitRepoPlugins, BNSetBundledPluginDi
 ///
 /// You can instead call this through [`Session`] or [`script_helper`]
 pub fn init() {
+    match crate::product().as_str() {
+        "Binary Ninja Enterprise Client" | "Binary Ninja Ultimate" => {
+            crate::enterprise::checkout_license(Duration::from_secs(900)).expect("Failed to checkout license");
+        },
+        _ => {}
+    }
+    
     unsafe {
         let path = binja_path().join("plugins").into_os_string();
         let path = path.into_string().unwrap();
@@ -96,6 +104,13 @@ pub fn init() {
 ///
 /// ⚠️ Important! Must be called at the end of scripts. ⚠️
 pub fn shutdown() {
+    match crate::product().as_str() {
+        "Binary Ninja Enterprise Client" | "Binary Ninja Ultimate" => {
+            crate::enterprise::release_license()
+        },
+        _ => {}
+    }
+    
     unsafe { binaryninjacore_sys::BNShutdown() };
 }