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

---
 rust/src/enterprise.rs | 66 ++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 66 insertions(+)

(limited to 'rust/src/enterprise.rs')

diff --git a/rust/src/enterprise.rs b/rust/src/enterprise.rs
index 232526dc..5b8877f3 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()) }
 }
-- 
cgit v1.3.1