extract run_command into package

commands/run_command.go

@@ -0,0 +1,120 @@
+package commands
+
+import (
+	"fmt"
+	"os"
+
+	"golang.org/x/crypto/ssh"
+)
+
+func ResolvePort(host string) string {
+
+	if len(host) > 0 {
+		hasPort := false
+		for i := len(host) - 1; i >= 0; i-- {
+			if host[i] == ':' {
+				hasPort = true
+				break
+			}
+			if host[i] == ']' {
+				break
+			}
+		}
+		if !hasPort {
+			host = host + ":22"
+		}
+	}
+
+	return host
+
+}
+
+func GetSigner(private_key_file string) ssh.Signer {
+
+	key_bytes, err := os.ReadFile(private_key_file)
+	if err != nil {
+		fmt.Fprintf(os.Stderr, "Error reading private key: %v\n", err)
+		os.Exit(1)
+	}
+
+	signer, err := ssh.ParsePrivateKey(key_bytes)
+	if err != nil {
+		fmt.Fprintf(os.Stderr, "Error parsing private key: %v\n", err)
+		os.Exit(1)
+	}
+
+	return signer
+
+}
+
+func GetClient(private_key_file, username, host string) *ssh.Client {
+
+	signer := GetSigner(private_key_file)
+
+	config := &ssh.ClientConfig{
+		User: username,
+		Auth: []ssh.AuthMethod{
+			ssh.PublicKeys(signer),
+		},
+		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
+	}
+
+	addr := ResolvePort(host)
+
+	client, err := ssh.Dial("tcp", addr, config)
+	if err != nil {
+		fmt.Fprintf(os.Stderr, "Error connecting to server: %v\n", err)
+		os.Exit(1)
+	}
+
+	return client
+}
+
+type RunCommandArgs struct {
+	Command string `json:"command"`
+	RemoteCommand string `json:"remote_command"`
+	PrivateKeyFile string `json:"private_key_file"`
+	Username string `json:"username"`
+	Host string `json:"host"`
+}
+
+func run_command_args() RunCommandArgs {
+
+	if len(os.Args) != 6 {
+		fmt.Fprintf(os.Stderr, "Usage: %s run_command <remote-command> <private-key> <username> <host>\n", os.Args[0])
+		os.Exit(1)
+	}
+
+	return RunCommandArgs{
+		os.Args[1],
+		os.Args[2],
+		os.Args[3],
+		os.Args[4],
+		os.Args[5],
+	}
+
+}
+
+func RunCommand() {
+
+	args := run_command_args()
+
+	client := GetClient(args.PrivateKeyFile, args.Username, args.Host)
+	defer client.Close()
+
+	session, err := client.NewSession()
+	if err != nil {
+		fmt.Fprintf(os.Stderr, "Error creating SSH session: %v\n", err)
+		os.Exit(1)
+	}
+	defer session.Close()
+
+	session.Stdout = os.Stdout
+	session.Stderr = os.Stderr
+	if err := session.Run(args.RemoteCommand); err != nil {
+		fmt.Fprintf(os.Stderr, "Error running command: %v\n", err)
+		os.Exit(1)
+	}
+	return
+
+}

main.go

@@ -7,6 +7,7 @@ import (
 	"path/filepath"
 
 	"golang.org/x/crypto/ssh"
+	"periodic/prddeploy/commands"
 )
 
 type UploadArgs struct {
@@ -43,6 +44,7 @@ func upload_read_args() UploadArgs {
 	}
 }
 
+/*
 func run_command_args() RunCommandArgs {
 
 	if len(os.Args) != 6 {
@@ -59,6 +61,7 @@ func run_command_args() RunCommandArgs {
 	}
 
 }
+*/
 
 func get_signer(private_key_file string) ssh.Signer {
 
@@ -176,6 +179,7 @@ func upload_file() {
 	fmt.Println("Success")
 }
 
+/*
 func run_command() {
 
 	args := run_command_args()
@@ -199,13 +203,20 @@ func run_command() {
 	return
 
 }
+*/
 
 func main() {
+
+	if len(os.Args) < 2 {
+		fmt.Fprintf(os.Stderr, "Usage: %s <command> [args...]", os.Args[0])
+		os.Exit(1)
+	}
+
 	switch os.Args[1] {
 	case "upload_file":
 		upload_file()
 	case "run_command":
-		run_command()
+		commands.RunCommand()
 	default: 
 		fmt.Println("HERE")
 	}