From 03e7d236919f5ea8a32162147da2b09a813e7b6b Mon Sep 17 00:00:00 2001
From: Shreerang Kale <shreerangkale@gmail.com>
Date: Thu, 12 Jun 2025 18:02:43 +0530
Subject: [PATCH] Add env variable to set tmkms mode

---
 .../compose/docker-compose-tmkms.yml             |  1 +
 stack-orchestrator/config/tmkms/run.sh           | 16 +++++++++++-----
 .../container-build/cerc-tmkms/build.sh          |  2 +-
 3 files changed, 13 insertions(+), 6 deletions(-)

diff --git a/stack-orchestrator/compose/docker-compose-tmkms.yml b/stack-orchestrator/compose/docker-compose-tmkms.yml
index 8bc5a62..42de778 100644
--- a/stack-orchestrator/compose/docker-compose-tmkms.yml
+++ b/stack-orchestrator/compose/docker-compose-tmkms.yml
@@ -8,6 +8,7 @@ services:
       NODE_IP: ${NODE_IP}
       NODE_PORT: ${NODE_PORT:-26659}
       KEY_PREFIX: ${KEY_PREFIX}
+      TMKMS_MODE: ${TMKMS_MODE:-softsign}
     volumes:
       - tmkms-data:/root/tmkms
       - ../config/tmkms/run.sh:/opt/run.sh
diff --git a/stack-orchestrator/config/tmkms/run.sh b/stack-orchestrator/config/tmkms/run.sh
index a8c67f5..107270f 100755
--- a/stack-orchestrator/config/tmkms/run.sh
+++ b/stack-orchestrator/config/tmkms/run.sh
@@ -12,11 +12,12 @@ TMKMS_SECRETS_DIR=$TMKMS_HOME/secrets
 TMKMS_STATE_DIR=$TMKMS_HOME/state
 
 # Initialize tmkms config if priv_validator_key does not exist
-if [[ ! -f "$TMKMS_SECRETS_DIR/priv_validator_key" ]]; then
+if [[ ! -f "$TMKMS_HOME/tmkms.toml" ]]; then
   echo "Initializing tmkms configuration..."
 
-  # TODO: run tmkms yubihsm setup
   tmkms init $TMKMS_HOME
+else
+  echo "tmkms configuration already exists. Skipping initialization."
 fi
 
 # Configure tmkms.toml and handle key import/copy based on TMKMS_MODE
@@ -43,12 +44,15 @@ case "$TMKMS_MODE" in
     auth = { key = 1, password = "password" }
 EOF
 
-    # Import the private validator key into tmkms for yubihsm (only if not already present)
+    # Setup YubiHSM
+    tmkms yubihsm setup -c $TMKMS_HOME/tmkms.toml
 
-    # TODO: Check yubihsm keys list
-    if [[ ! -f "$TMKMS_SECRETS_DIR/priv_validator_key" ]]; then
+    # Import the private validator key into tmkms for yubihsm (only if not already present)
+    if ! tmkms yubihsm keys list | grep -q "0x0001:"; then
       echo "Importing private validator key into tmkms for yubihsm..."
       tmkms yubihsm keys import -i 1 $INPUT_PRIV_KEY_FILE -c $TMKMS_HOME/tmkms.toml
+    else
+      echo "Key 0x0001 already present in YubiHSM. Skipping import."
     fi
     ;;
 
@@ -78,6 +82,8 @@ EOF
     if [[ ! -f "$TMKMS_SECRETS_DIR/priv_validator_key" ]]; then
       echo "Importing private validator key into tmkms for softsign..."
       tmkms softsign import $INPUT_PRIV_KEY_FILE $TMKMS_SECRETS_DIR/priv_validator_key
+    else
+      echo "Softsign key already present. Skipping import."
     fi
     ;;
 
diff --git a/stack-orchestrator/container-build/cerc-tmkms/build.sh b/stack-orchestrator/container-build/cerc-tmkms/build.sh
index 8a8da61..918cceb 100755
--- a/stack-orchestrator/container-build/cerc-tmkms/build.sh
+++ b/stack-orchestrator/container-build/cerc-tmkms/build.sh
@@ -7,4 +7,4 @@ source ${CERC_CONTAINER_BASE_DIR}/build-base.sh
 SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )
 
 # TODO: Use BACKEND=yubihsm build command arg
-docker build -t cerc/tmkms:local --build-arg BACKEND=yubihsm ${build_command_args} -f ${SCRIPT_DIR}/Dockerfile ${SCRIPT_DIR}
+docker build -t cerc/tmkms:local ${build_command_args} -f ${SCRIPT_DIR}/Dockerfile ${SCRIPT_DIR}