Refactor SSH command execution and enhance port checking logic in Lua script

scripts/auto-boot-ollama-host.lua

@@ -12,6 +12,26 @@ local function getenv(name, def)
   return (v ~= nil and v ~= "") and v or def
 end
 
+-- Check if a TCP port is accepting connections within a timeout (seconds)
+local function port_is_up(host, port, timeout_sec)
+  host = tostring(host or "127.0.0.1")
+  port = tonumber(port or 0) or 0
+  local timeout = tonumber(timeout_sec or 1) or 1
+  if port <= 0 then return false end
+
+  local deadline = socket.gettime() + timeout
+  while socket.gettime() < deadline do
+    local tcp = socket.tcp()
+    if not tcp then return false end
+    tcp:settimeout(1)
+    local ok = tcp:connect(host, port)
+    tcp:close()
+    if ok then return true end
+    socket.sleep(0.5)
+  end
+  return false
+end
+
 -- ---- Config via env ----
 local CONTAINER_NAME 		= getenv("CONTAINER_NAME", "paperless-ai")
 local SINCE          		= getenv("SINCE", "0s")
@@ -19,8 +39,6 @@ local OLLAMA_HOST    		= getenv("OLLAMA_HOST", "192.168.222.12")
 local OLLAMA_PORT    		= tonumber(getenv("OLLAMA_PORT", "11434"))
 local SSH_PORT		   		= tonumber(getenv("SSH_PORT", "22"))
 local SSH_IDENTITY_FILE = getenv("SSH_IDENTITY_FILE", "/root/.ssh/id")  -- e.g. "/path/to/id_rsa"
-local SSH_PRIVATE_KEY  	= getenv("SSH_PRIVATE_KEY", "")
-local SSH_PUBLIC_KEY  	= getenv("SSH_PUBLIC_KEY", "")
 local ERROR_PATTERN  		= getenv(
   "ERROR_PATTERN",
   ("connect EHOSTUNREACH %s:%d"):format(OLLAMA_HOST, OLLAMA_PORT)
@@ -65,57 +83,65 @@ local function send_wol(mac_str, bcast_ip, port)
   return ok ~= nil, err
 end
 
--- Kept for reference; not used to keep parity with your minimal bash
--- local function port_is_up(host, port, timeout_sec)
---   local deadline = socket.gettime() + (timeout_sec or 1)
---   repeat
---     local tcp = socket.tcp(); tcp:settimeout(1)
---     if tcp:connect(host, port) then tcp:close(); return true end
---     tcp:close(); socket.sleep(0.5)
---   until socket.gettime() >= deadline
---   return false
--- end
-
---local function createSSHKeyFilesFromEnv()
---  if SSH_PRIVATE_KEY == "" or SSH_PUBLIC_KEY == "" then
---    log("SSH_PRIVATE_KEY or SSH_PUBLIC_KEY env var is empty, skipping SSH key file creation.")
---    return
---  end
---
---  -- Ensure .ssh directory exists
---  local ssh_dir = SSH_IDENTITY_FILE:match("^(.*)/[^/]+$")
---  if ssh_dir then
---    os.execute(("mkdir -p %q && chmod 700 %q"):format(ssh_dir, ssh_dir))
---  end
---
---  local priv_fh = io.open(SSH_IDENTITY_FILE, "w")
---  if not priv_fh then
---    log("Failed to open SSH identity file for writing: " .. SSH_IDENTITY_FILE)
---    return
---  end
---  priv_fh:write(SSH_PRIVATE_KEY)
---  priv_fh:close()
---  os.execute(("chmod 600 %q"):format(SSH_IDENTITY_FILE))
---  log("Wrote SSH private key to " .. SSH_IDENTITY_FILE)
---
---  local pub_fh = io.open(SSH_IDENTITY_FILE .. ".pub", "w")
---  if not pub_fh then
---    log("Failed to open SSH public key file for writing: " .. SSH_IDENTITY_FILE .. ".pub")
---    return
---  end
---  pub_fh:write(SSH_PUBLIC_KEY)
---  pub_fh:close()
---  os.execute(("chmod 644 %q"):format(SSH_IDENTITY_FILE .. ".pub"))
---  log("Wrote SSH public key to " .. SSH_IDENTITY_FILE .. ".pub")
---
---  -- Unset the env vars for security
---  os.setenv("SSH_PRIVATE_KEY", "")
---  os.setenv("SSH_PUBLIC_KEY", "")
---end
+-- Execute a remote command over SSH.
+-- Signature must remain: ssh(command, user, host, port, identity_file)
+local function ssh(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
 
-local function main()
-  -- createSSHKeyFilesFromEnv()
+  -- 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=10",
+    "-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)
+
+  -- Quote remote command so the local shell treats it as a single arg
+  table.insert(pieces, sq(command))
 
+  -- 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("SSH exec: " .. full)
+  local ok, reason, code = os.execute(full)
+  if ok == true or ok == 0 then
+    log("SSH command completed successfully")
+    return true
+  else
+    local msg = string.format("SSH failed: reason=%s code=%s", tostring(reason), tostring(code))
+    log(msg)
+    return false, msg
+  end
+end
+
+
+local function main()
   log(("Watching container='%s' since='%s'"):format(CONTAINER_NAME, SINCE))
   log(("Looking for pattern: %q"):format(ERROR_PATTERN))
 
@@ -146,7 +172,7 @@ local function main()
 
 			if port_is_up(OLLAMA_HOST, SSH_PORT, 60) then
 				log("SSH is reachable. Starting ollama service...")
-				--ssh("wsl.exe -d Debian -- sudo systemctl enable --now ollama && sudo systemctl start ollama", "micro", OLLAMA_HOST, SSH_PORT, SSH_IDENTITY_FILE)
+				ssh("wsl.exe -d Debian -- sudo systemctl enable --now ollama && sudo systemctl start ollama", "micro", OLLAMA_HOST, SSH_PORT, SSH_IDENTITY_FILE)
       end
     end
   end