auto-boot-ollama-host/scripts/ssh.lua

-- SSH utilities module for auto-boot-ollama-host
-- Provides SSH command execution functionality

local utils = require("utils")

local ssh_module = {}

-- Check if DEBUG environment variable is set to "true"
local function is_debug()
  local debug_env = os.getenv("DEBUG")
  return debug_env and (string.lower(debug_env) == "true" or debug_env == "1")
end

-- Execute a remote command over SSH
-- Signature: ssh(command, user, host, port, identity_file)
function ssh_module.execute(command, user, host, port, identity_file)
  -- Basic validation and defaults
  user = tostring(user or "")
  host = tostring(host or "")
  port = tonumber(port or 22) or 22
  identity_file = tostring(identity_file or "")

  -- Quote a string for safe single-quoted POSIX shell context
  local function sq(s)
    -- Replace ' with: '\''  (close, escape quote, reopen)
    return "'" .. tostring(s):gsub("'", "'\\''") .. "'"
  end

  -- Build base ssh command (run locally)
  -- -oBatchMode to avoid interactive prompts
  -- -oConnectTimeout for faster failure
  -- -oStrictHostKeyChecking uses known_hosts; adjust if needed
  local dest = (user ~= "" and (user .. "@" .. host) or host)
  local pieces = {
    "ssh",
    "-p", tostring(port),
    "-o", "BatchMode=yes",
    "-o", "ConnectTimeout=30",
    "-o", "ServerAliveInterval=5",
    "-o", "ServerAliveCountMax=1",
    "-o", "UserKnownHostsFile=/root/.ssh/known_hosts",
    "-o", "StrictHostKeyChecking=yes",
  }

  if identity_file ~= "" then
    table.insert(pieces, "-i")
    table.insert(pieces, identity_file)
  end
  table.insert(pieces, dest)

  -- Pass remote command as provided; caller is responsible for proper quoting
  table.insert(pieces, "--")
  -- Quote the remote command to prevent shell interpretation of && and ||
  table.insert(pieces, "'" .. command:gsub("'", "'\\''") .. "'")

  -- Join with spaces for os.execute
  local function join(args)
    -- We only quote the remote command explicitly. Other args are simple tokens.
    return table.concat(args, " ")
  end

  local full = join(pieces)

  -- Log based on DEBUG environment variable
  if is_debug() then
    utils.log("SSH exec: " .. full)
  else
    utils.log("SSH exec: " .. "'" .. command:gsub("'", "'\\''") .. "'")
  end

  local ok, reason, code = os.execute(full)
  if ok == true or ok == 0 then
    utils.log("SSH command completed successfully")
    return true
  else
    local msg = string.format("SSH failed: reason=%s code=%s", tostring(reason), tostring(code))
    utils.log(msg)
    return false, msg
  end
end

-- Execute a remote command over SSH and return the output
-- Signature: ssh.execute_with_output(command, user, host, port, identity_file)
-- Returns: success, output, error_message
function ssh_module.execute_with_output(command, user, host, port, identity_file)
  -- Basic validation and defaults
  user = tostring(user or "")
  host = tostring(host or "")
  port = tonumber(port or 22) or 22
  identity_file = tostring(identity_file or "")

  -- Build base ssh command (run locally)
  local dest = (user ~= "" and (user .. "@" .. host) or host)
  local pieces = {
    "ssh",
    "-p", tostring(port),
    "-o", "BatchMode=yes",
    "-o", "ConnectTimeout=30",
    "-o", "ServerAliveInterval=5",
    "-o", "ServerAliveCountMax=1",
    "-o", "UserKnownHostsFile=/root/.ssh/known_hosts",
    "-o", "StrictHostKeyChecking=yes",
  }

  if identity_file ~= "" then
    table.insert(pieces, "-i")
    table.insert(pieces, identity_file)
  end
  table.insert(pieces, dest)

  -- Pass remote command as provided
  table.insert(pieces, "--")
  -- Quote the remote command to prevent shell interpretation of && and ||
  table.insert(pieces, "'" .. command:gsub("'", "'\\''") .. "'")

  -- Join with spaces for io.popen
  local function join(args)
    return table.concat(args, " ")
  end

  local full = join(pieces)

  -- Log based on DEBUG environment variable
  if is_debug() then
    utils.log("SSH exec (with output): " .. full)
  else
    utils.log("SSH exec (with output): " .. "'" .. command:gsub("'", "'\\''") .. "'")
  end

  -- Use io.popen to capture output
  local fh = io.popen(full, "r")
  if not fh then
    return false, "", "Failed to open SSH command"
  end

  local output = fh:read("*a")
  local success, reason, code = fh:close()

  if success then
    utils.log("SSH command completed successfully with output")
    return true, output, nil
  else
    local msg = string.format("SSH failed: reason=%s code=%s", tostring(reason), tostring(code))
    utils.log(msg)
    return false, output, msg
  end
end

return ssh_module