From ac568a5153db8b4bd2472389c9e3c447a48cae61 Mon Sep 17 00:00:00 2001
From: Shreerang Kale <shreerangkale@gmail.com>
Date: Thu, 12 Jun 2025 15:16:53 +0530
Subject: [PATCH]  Update script to use yubihsm

---
 stack-orchestrator/config/tmkms/run.sh        | 107 ++++++++++++------
 .../container-build/cerc-tmkms/build.sh       |   2 +-
 2 files changed, 73 insertions(+), 36 deletions(-)

diff --git a/stack-orchestrator/config/tmkms/run.sh b/stack-orchestrator/config/tmkms/run.sh
index d317944..14ed884 100755
--- a/stack-orchestrator/config/tmkms/run.sh
+++ b/stack-orchestrator/config/tmkms/run.sh
@@ -11,50 +11,87 @@ INPUT_PRIV_KEY_FILE=$TMKMS_HOME/tmp/priv_validator_key.json
 TMKMS_SECRETS_DIR=$TMKMS_HOME/secrets
 TMKMS_STATE_DIR=$TMKMS_HOME/state
 
-# Check if priv_validator_key in SECRETS_DIR exists
+# Initialize tmkms config if priv_validator_key does not exist
 if [[ ! -f "$TMKMS_SECRETS_DIR/priv_validator_key" ]]; then
-  # Initialize tmkms config
   echo "Initializing tmkms configuration..."
+
+  # TODO: run tmkms yubihsm setup
   tmkms init $TMKMS_HOME
-
-  # Import the private validator key into tmkms
-  echo "Importing private validator key into tmkms..."
-  tmkms softsign import $INPUT_PRIV_KEY_FILE $TMKMS_SECRETS_DIR/priv_validator_key
-
-  # Remove the original private validator key
-  rm -rf $INPUT_PRIV_KEY_FILE
-
-else
-  echo "tmkms configuration already exists. Skipping initialization and cleaning up any existing input private validator key files..."
-  # Remove the original private validator key as it is not needed
-  if [[ -f "$INPUT_PRIV_KEY_FILE" ]]; then
-    rm -rf $INPUT_PRIV_KEY_FILE
-  fi
 fi
 
-# Update tmkms.toml
-echo "Updating tmkms.toml with chain_id, node IP, and key prefixes..."
+# Configure tmkms.toml and handle key import/copy based on TMKMS_MODE
+case "$TMKMS_MODE" in
+  "yubihsm")
+    # Add chain configuration for yubihsm
+    # TODO: Take password from env var
+    cat <<EOF > $TMKMS_HOME/tmkms.toml
 
-# Add chain configuration
-cat <<EOF > $TMKMS_HOME/tmkms.toml
+    [[chain]]
+    id = "$CHAIN_ID"
+    key_format = { type = "cosmos-json", account_key_prefix = "${KEY_PREFIX}pub", consensus_key_prefix = "${KEY_PREFIX}valconspub" }
+    state_file = "$TMKMS_STATE_DIR/priv_validator_state.json"
 
-  [[chain]]
-  id = "$CHAIN_ID"
-  key_format = { type = "cosmos-json", account_key_prefix = "${KEY_PREFIX}pub", consensus_key_prefix = "${KEY_PREFIX}valconspub" }
-  state_file = "$TMKMS_STATE_DIR/priv_validator_state.json"
+    [[validator]]
+    chain_id = "$CHAIN_ID"
+    addr = "tcp://$NODE_IP:$NODE_PORT"
+    secret_key = "$TMKMS_SECRETS_DIR/kms-identity.key"
+    protocol_version = "v0.34"
+    reconnect = true
 
-  [[validator]]
-  chain_id = "$CHAIN_ID"
-  addr = "tcp://$NODE_IP:$NODE_PORT"
-  secret_key = "$TMKMS_SECRETS_DIR/kms-identity.key"
-  protocol_version = "v0.34"
-  reconnect = true
-
-  [[providers.softsign]]
-  key_type = "consensus"
-  path = "$TMKMS_SECRETS_DIR/priv_validator_key"
-  chain_ids = ["$CHAIN_ID"]
+    [[providers.yubihsm]]
+    adapter = { type = "usb" }
+    auth = { key = 1, password = "password" }
 EOF
 
+    # Import the private validator key into tmkms for yubihsm (only if not already present)
+
+    # TODO: Check yubihsm keys list
+    if [[ ! -f "$TMKMS_SECRETS_DIR/priv_validator_key" ]]; 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
+    fi
+    ;;
+
+  "softsign")
+    # Add chain configuration for softsign
+    cat <<EOF > $TMKMS_HOME/tmkms.toml
+
+    [[chain]]
+    id = "$CHAIN_ID"
+    key_format = { type = "cosmos-json", account_key_prefix = "${KEY_PREFIX}pub", consensus_key_prefix = "${KEY_PREFIX}valconspub" }
+    state_file = "$TMKMS_STATE_DIR/priv_validator_state.json"
+
+    [[validator]]
+    chain_id = "$CHAIN_ID"
+    addr = "tcp://$NODE_IP:$NODE_PORT"
+    secret_key = "$TMKMS_SECRETS_DIR/kms-identity.key"
+    protocol_version = "v0.34"
+    reconnect = true
+
+    [[providers.softsign]]
+    key_type = "consensus"
+    path = "$TMKMS_SECRETS_DIR/priv_validator_key"
+    chain_ids = ["$CHAIN_ID"]
+EOF
+
+    # Import the private validator key into tmkms for softsign (only if not already present)
+    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
+    fi
+    ;;
+
+  *)
+    echo "Error: TMKMS_MODE environment variable not set or invalid. Please set it to 'yubihsm' or 'softsign'."
+    exit 1
+    ;;
+esac
+
+# Remove the original input private validator key file after processing
+if [[ -f "$INPUT_PRIV_KEY_FILE" ]]; then
+  rm -rf $INPUT_PRIV_KEY_FILE
+fi
+
+# Start tmkms
 echo "Starting tmkms..."
 tmkms start --config $TMKMS_HOME/tmkms.toml
diff --git a/stack-orchestrator/container-build/cerc-tmkms/build.sh b/stack-orchestrator/container-build/cerc-tmkms/build.sh
index 918cceb..8a8da61 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_command_args} -f ${SCRIPT_DIR}/Dockerfile ${SCRIPT_DIR}
+docker build -t cerc/tmkms:local --build-arg BACKEND=yubihsm ${build_command_args} -f ${SCRIPT_DIR}/Dockerfile ${SCRIPT_DIR}