BadUSB: Added support for constants and variables definition

applications/main/bad_usb/helpers/ducky_script.c

@@ -132,6 +132,114 @@ int32_t ducky_error(BadUsbScript* bad_usb, const char* text, ...) {
     return SCRIPT_STATE_ERROR;
 }
 
+const char* ducky_map_get(map_str_t map, const char* key) {
+    if(key[0] == '$' || key[0] == '#') key++;
+    // get the value from the map
+    const char** value = map_str_get(map, key);
+    return value ? *value : NULL;
+}
+
+void ducky_maps_free(BadUsbScript* bad_usb, bool init) {
+    // hashmaps must be cleared when script is stopped/paused otherwise we could have a
+    // multiple #constant definition error
+    map_str_clear(bad_usb->variables);
+    map_str_clear(bad_usb->constants);
+    map_str_clear(bad_usb->constants_sharp);
+
+    // hashmaps must be re-initialized if script has been stopped/paused to avoid app crash
+    // hashmap must not be re-initialized if script is closed
+    if (init) {
+        map_str_init(bad_usb->variables);
+        map_str_init(bad_usb->constants);
+        map_str_init(bad_usb->constants_sharp);
+    }
+}
+
+int32_t ducky_define(BadUsbScript* bad_usb, const char* param, bool is_constant) {
+
+    char* space_pos;
+    // Variables must be assigned using =
+    if(!is_constant)
+        if (strchr(param, '=') == NULL)
+            return ducky_error(bad_usb, "No valid assignment for variable", param);
+        else {
+            // get variable name and value splitting the string
+            space_pos = strchr(param, '=');
+        }
+    else {
+        // get variable name and value splitting the string
+        space_pos = strchr(param, ' ');
+    }
+
+    if(space_pos == NULL) return ducky_error(bad_usb, "No valid name for variable or constant", param);
+
+
+    size_t var_name_length = space_pos - param;
+    // Remove ending spaces from the variable name
+    while(var_name_length > 0 && param[var_name_length - 1] == ' ')
+        var_name_length--;
+
+    // Memory allocation
+    char* var_name = (char*)malloc(var_name_length + 1);
+    if (!var_name) {
+        free(var_name); // free var_name before exit
+        return ducky_error(bad_usb, "Var name NULL", param);
+    }
+
+    // Copy variable name and add the null terminator
+    memcpy(var_name, param, var_name_length);
+    var_name[var_name_length] = '\0';
+
+
+    // Check if the variable starts with a '$' and remove it
+    if(is_constant) {
+        if (var_name[0] == '$')
+            return ducky_error(bad_usb, "Constant cannot start with $", var_name);
+        else if(var_name[0] == '#' && ducky_map_get(bad_usb->constants_sharp, var_name) != NULL)
+            return ducky_error(bad_usb, "Multiple #constant definition for %s", var_name);
+        else if(var_name[0] != '#' && ducky_map_get(bad_usb->constants, var_name) != NULL)
+            return ducky_error(bad_usb, "Multiple constant definition for %s", var_name);
+    }
+    else {
+        if (var_name[0] != '$')
+            return ducky_error(bad_usb, "Variable must start with $", var_name);
+    }
+
+    // Memory allocation
+    size_t var_value_length = strlen(space_pos + 1);
+
+    // Remove starting spaces from the variable value
+    while(*(space_pos+1) == ' ') {
+        space_pos++;
+        var_value_length--;
+    }
+
+    char* var_value = (char*)malloc(var_value_length + 1);
+    if (!var_value) {
+        free(var_name); // free var_name before exit
+        free(var_value); // free var_name before exit
+        return ducky_error(bad_usb, "Var value is NULL", param);
+    }
+
+    // Copy variable value and add the null terminator
+    memcpy(var_value, space_pos + 1, var_value_length);
+    var_value[var_value_length] = '\0';
+
+    if(var_name == NULL || var_value == NULL)
+        return ducky_error(bad_usb, "Var name or value is NULL", param);
+
+    if (is_constant)
+        if (var_name[0] == '#')
+            map_str_set_at(bad_usb->constants_sharp, ++var_name, var_value);
+        else
+            map_str_set_at(bad_usb->constants, var_name, var_value);
+    else
+        map_str_set_at(bad_usb->variables, ++var_name, var_value);
+    FURI_LOG_D(WORKER_TAG, "Var Name: %s - Var Value: %s", var_name, var_value);
+
+    return 0;
+}
+
 bool ducky_string(BadUsbScript* bad_usb, const char* param) {
     uint32_t i = 0;
 
@@ -507,9 +615,11 @@ static int32_t bad_usb_worker(void* context) {
                 } else if(flags & WorkerEvtStartStop) {
                     worker_state = BadUsbStateIdle; // Stop executing script
                     bad_usb->hid->release_all(bad_usb->hid_inst);
+                    ducky_maps_free(bad_usb, true);
                 } else if(flags & WorkerEvtDisconnect) {
                     worker_state = BadUsbStateNotConnected; // USB disconnected
                     bad_usb->hid->release_all(bad_usb->hid_inst);
+                    ducky_maps_free(bad_usb, true);
                 } else if(flags & WorkerEvtPauseResume) {
                     pause_state = BadUsbStateRunning;
                     worker_state = BadUsbStatePaused; // Pause
@@ -530,11 +640,13 @@ static int32_t bad_usb_worker(void* context) {
                     worker_state = BadUsbStateScriptError;
                     bad_usb->st.state = worker_state;
                     bad_usb->hid->release_all(bad_usb->hid_inst);
+                    ducky_maps_free(bad_usb, true);
                 } else if(delay_val == SCRIPT_STATE_END) { // End of script
                     delay_val = 0;
                     worker_state = BadUsbStateIdle;
                     bad_usb->st.state = BadUsbStateDone;
                     bad_usb->hid->release_all(bad_usb->hid_inst);
+                    ducky_maps_free(bad_usb, true);
                     continue;
                 } else if(delay_val == SCRIPT_STATE_STRING_START) { // Start printing string with delays
                     delay_val = bad_usb->defdelay;
@@ -563,6 +675,7 @@ static int32_t bad_usb_worker(void* context) {
                 } else if(flags & WorkerEvtDisconnect) {
                     worker_state = BadUsbStateNotConnected; // USB disconnected
                     bad_usb->hid->release_all(bad_usb->hid_inst);
+                    ducky_maps_free(bad_usb, true);
                 }
                 bad_usb->st.state = worker_state;
                 continue;
@@ -578,10 +691,12 @@ static int32_t bad_usb_worker(void* context) {
                     worker_state = BadUsbStateIdle; // Stop executing script
                     bad_usb->st.state = worker_state;
                     bad_usb->hid->release_all(bad_usb->hid_inst);
+                    ducky_maps_free(bad_usb, true);
                 } else if(flags & WorkerEvtDisconnect) {
                     worker_state = BadUsbStateNotConnected; // USB disconnected
                     bad_usb->st.state = worker_state;
                     bad_usb->hid->release_all(bad_usb->hid_inst);
+                    ducky_maps_free(bad_usb, true);
                 } else if(flags & WorkerEvtPauseResume) {
                     if(pause_state == BadUsbStateRunning) {
                         if(delay_val > 0) {
@@ -612,9 +727,11 @@ static int32_t bad_usb_worker(void* context) {
                 } else if(flags & WorkerEvtStartStop) {
                     worker_state = BadUsbStateIdle; // Stop executing script
                     bad_usb->hid->release_all(bad_usb->hid_inst);
+                    ducky_maps_free(bad_usb, true);
                 } else if(flags & WorkerEvtDisconnect) {
                     worker_state = BadUsbStateNotConnected; // USB disconnected
                     bad_usb->hid->release_all(bad_usb->hid_inst);
+                    ducky_maps_free(bad_usb, true);
                 } else if(flags & WorkerEvtPauseResume) {
                     pause_state = BadUsbStateStringDelay;
                     worker_state = BadUsbStatePaused; // Pause
@@ -637,7 +754,7 @@ static int32_t bad_usb_worker(void* context) {
             (worker_state == BadUsbStateScriptError)) { // State: error
             uint32_t flags =
                 bad_usb_flags_get(WorkerEvtEnd, FuriWaitForever); // Waiting for exit command
-
+            ducky_maps_free(bad_usb, true);
             if(flags & WorkerEvtEnd) {
                 break;
             }
@@ -678,6 +795,11 @@ BadUsbScript* bad_usb_script_open(FuriString* file_path, BadUsbHidInterface inte
 
     bad_usb->thread = furi_thread_alloc_ex("BadUsbWorker", 2048, bad_usb_worker, bad_usb);
     furi_thread_start(bad_usb->thread);
+
+    map_str_init(bad_usb->variables);
+    map_str_init(bad_usb->constants);
+    map_str_init(bad_usb->constants_sharp);
+
     return bad_usb;
 } //-V773
 
@@ -687,6 +809,7 @@ void bad_usb_script_close(BadUsbScript* bad_usb) {
     furi_thread_join(bad_usb->thread);
     furi_thread_free(bad_usb->thread);
     furi_string_free(bad_usb->file_path);
+    ducky_maps_free(bad_usb, false);
     free(bad_usb);
 }
 

applications/main/bad_usb/helpers/ducky_script_commands.c

@@ -57,9 +57,79 @@ static int32_t ducky_fnc_defstrdelay(BadUsbScript* bad_usb, const char* line, in
     return 0;
 }
 
+static int32_t ducky_fnc_define(BadUsbScript* bad_usb, const char* line, int32_t param) {
+    line = &line[ducky_get_command_len(line) + 1];
+    return ducky_define(bad_usb, line, param);
+}
+
+// Function to check if a character is valid in a variable name
+int is_valid_var_char(char c, int is_first) {
+    if (is_first) {
+        return isalpha(c) || c == '_'; // First character must be a letter or underscore
+    }
+    return isalnum(c) || c == '_'; // Subsequent characters can be letters, digits, or underscores
+}
+
 static int32_t ducky_fnc_string(BadUsbScript* bad_usb, const char* line, int32_t param) {
     line = &line[ducky_get_command_len(line) + 1];
-    furi_string_set_str(bad_usb->string_print, line);
+
+    char* line_ptr = (char*)line;
+    bool first = true;
+
+    while (*line_ptr != '\0') {
+        // Skip leading spaces safely
+        while (*line_ptr != '\0' && *line_ptr == ' ') line_ptr++;
+        if (*line_ptr == '\0') break;
+
+        // Find the next space
+        char* end = strchr(line_ptr, ' ');
+        size_t token_length = (end) ? (size_t)(end - line_ptr) : strlen(line_ptr);
+
+        // Allocate token buffer dynamically
+        char* token = (char*)malloc(token_length + 1);
+        if (!token) break; // Handle allocation failure
+        strncpy(token, line_ptr, token_length);
+        token[token_length] = '\0';
+
+        // Determine the appropriate map
+        map_str_t* temp_map = &(bad_usb->constants);
+        if (token[0] == '#')
+            temp_map = &(bad_usb->constants_sharp);
+        else if (token[0] == '$')
+            temp_map = &(bad_usb->variables);
+
+        // Handle mapping
+        if (first) {
+            if (ducky_map_get(*temp_map, token) == NULL)
+                furi_string_set_str(bad_usb->string_print, token);
+            else
+                furi_string_set_str(bad_usb->string_print, ducky_map_get(*temp_map, token));
+            first = false;
+        } else {
+            if (ducky_map_get(*temp_map, token) == NULL)
+                furi_string_cat_str(bad_usb->string_print, token);
+            else
+                furi_string_cat_str(bad_usb->string_print, ducky_map_get(*temp_map, token));
+        }
+
+        // Free allocated memory
+        free(token);
+
+        // Add spaces if necessary
+        if (end) {
+            size_t space_count = 0;
+            while (*end == ' ') {
+                space_count++;
+                end++;
+            }
+            for (size_t i = 0; i < space_count; ++i)
+                furi_string_cat_str(bad_usb->string_print, " ");
+            line_ptr = end;
+        } else {
+            break; // Prevents undefined behavior if end is NULL
+        }
+    }
+
     if(param == 1) {
         furi_string_cat(bad_usb->string_print, "\n");
     }
@@ -257,6 +327,8 @@ static const DuckyCmd ducky_commands[] = {
     {"REM", NULL, -1},
     {"ID", NULL, -1},
     {"DELAY", ducky_fnc_delay, -1},
+    {"DEFINE", ducky_fnc_define, 1},
+    {"VAR", ducky_fnc_define, 0},
     {"STRING", ducky_fnc_string, 0},
     {"STRINGLN", ducky_fnc_string, 1},
     {"DEFAULT_DELAY", ducky_fnc_defdelay, -1},

applications/main/bad_usb/helpers/ducky_script_i.h

@@ -6,9 +6,13 @@ extern "C" {
 
 #include <furi.h>
 #include <furi_hal.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <mlib/m-dict.h>
 #include "ducky_script.h"
 #include "bad_usb_hid.h"
 
+
 #define SCRIPT_STATE_ERROR        (-1)
 #define SCRIPT_STATE_END          (-2)
 #define SCRIPT_STATE_NEXT_LINE    (-3)
@@ -21,6 +25,8 @@ extern "C" {
 #define HID_MOUSE_INVALID 0
 #define HID_MOUSE_NONE    0
 
+DICT_DEF2(map_str, const char*, M_CSTR_OPLIST, const char*, M_CSTR_OPLIST)
+
 struct BadUsbScript {
     FuriHalUsbHidConfig hid_cfg;
     const BadUsbHidApi* hid;
@@ -46,6 +52,10 @@ struct BadUsbScript {
 
     FuriString* string_print;
     size_t string_print_pos;
+
+    map_str_t variables;
+    map_str_t constants;
+    map_str_t constants_sharp;
 };
 
 uint16_t ducky_get_keycode(BadUsbScript* bad_usb, const char* param, bool accept_chars);
@@ -74,10 +84,16 @@ bool ducky_altstring(BadUsbScript* bad_usb, const char* param);
 
 bool ducky_string(BadUsbScript* bad_usb, const char* param);
 
+int32_t ducky_define(BadUsbScript* bad_usb, const char* param, bool is_constant);
+
 int32_t ducky_execute_cmd(BadUsbScript* bad_usb, const char* line);
 
 int32_t ducky_error(BadUsbScript* bad_usb, const char* text, ...);
 
+const char* ducky_map_get(map_str_t map, const char* key);
+
+void ducky_maps_free(BadUsbScript* bad_usb, bool init);
+
 #ifdef __cplusplus
 }
 #endif