main.go

// ssh-inject - quick tool i made to stop typing passwords all the time
// got tired of manually copying ssh keys to servers so here we are

package main

import (
	"bufio"
	"fmt"
	"io"
	"net"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"strings"
	"time"

	"golang.org/x/crypto/ssh"
	"golang.org/x/term"
)

// ansi colors - nothing fancy just the basics
const (
	colorReset  = "\033[0m"
	colorRed    = "\033[31m"
	colorGreen  = "\033[32m"
	colorYellow = "\033[33m"
	colorBlue   = "\033[34m"
	colorPurple = "\033[35m"
	colorCyan   = "\033[36m"
	colorWhite  = "\033[37m"
)

// made this with some ascii generator online lol
var banner = `
` + colorCyan + `
   _____ _____ _    _   _____       _           _
  / ____/ ____| |  | | |_   _|     (_)         | |
 | (___| (___ | |__| |   | |  _ __  _  ___  ___| |_
  \___ \\___ \|  __  |   | | | '_ \| |/ _ \/ __| __|
  ____) |___) | |  | |  _| |_| | | | |  __/ (__| |_
 |_____/_____/|_|  |_| |_____|_| |_| |\___|\___|\__|
                                  _/ |
                                 |__/
` + colorReset + `
` + colorYellow + `  [ SSH Key Injection Tool v1.0 ]` + colorReset + `
` + colorPurple + `  [ Cross-Platform | No Dependencies ]` + colorReset + `
` + colorGreen + `  [ Host OS: %s | Arch: %s ]` + colorReset + `
`

var menuArt = `
` + colorCyan + `╔════════════════════════════════════════════════════════╗
║` + colorWhite + `                    MAIN MENU                           ` + colorCyan + `║
╠════════════════════════════════════════════════════════╣
║` + colorGreen + `  [1]` + colorWhite + ` Inject SSH Key to Remote Server                   ` + colorCyan + `║
║` + colorGreen + `  [2]` + colorWhite + ` Generate New SSH Key Pair                         ` + colorCyan + `║
║` + colorGreen + `  [3]` + colorWhite + ` View Current Public Key                           ` + colorCyan + `║
║` + colorGreen + `  [4]` + colorWhite + ` Test SSH Connection                               ` + colorCyan + `║
║` + colorRed + `  [5]` + colorWhite + ` Exit                                              ` + colorCyan + `║
╚════════════════════════════════════════════════════════╝` + colorReset + `
`

var osSelectArt = `
` + colorCyan + `╔════════════════════════════════════════════════════════╗
║` + colorWhite + `              SELECT REMOTE OS TYPE                     ` + colorCyan + `║
╠════════════════════════════════════════════════════════╣
║` + colorGreen + `  [1]` + colorWhite + ` Linux                                           ` + colorCyan + `║
║` + colorGreen + `  [2]` + colorWhite + ` macOS                                           ` + colorCyan + `║
║` + colorGreen + `  [3]` + colorWhite + ` Windows                                         ` + colorCyan + `║
╚════════════════════════════════════════════════════════╝` + colorReset + `
`

func main() {
	clearScreen()
	printBanner()

	reader := bufio.NewReader(os.Stdin)

	// main loop - just keep showing menu until they quit
	for {
		fmt.Println(menuArt)
		fmt.Print(colorYellow + "  Select option: " + colorReset)

		choice, _ := reader.ReadString('\n')
		choice = strings.TrimSpace(choice)

		switch choice {
		case "1":
			injectSSHKey(reader)
		case "2":
			generateSSHKey(reader)
		case "3":
			viewPublicKey()
		case "4":
			testConnection(reader)
		case "5":
			fmt.Println(colorGreen + "\n  Goodbye! Happy SSH-ing!" + colorReset)
			os.Exit(0)
		default:
			fmt.Println(colorRed + "\n  Invalid option. Please try again." + colorReset)
		}

		fmt.Print(colorYellow + "\n  Press Enter to continue..." + colorReset)
		reader.ReadString('\n')
		clearScreen()
		printBanner()
	}
}

// clearScreen - works on windows and unix, tested on my machines
func clearScreen() {
	switch runtime.GOOS {
	case "windows":
		cmd := exec.Command("cmd", "/c", "cls")
		cmd.Stdout = os.Stdout
		cmd.Run()
	default:
		// escape codes work on basically everything else
		fmt.Print("\033[H\033[2J")
	}
}

func printBanner() {
	fmt.Printf(banner, runtime.GOOS, runtime.GOARCH)
}

// helper funcs for getting ssh paths
// these should work on all os's since go handles the path separators

func getDefaultSSHDir() string {
	home, err := os.UserHomeDir()
	if err != nil {
		return ""
	}
	return filepath.Join(home, ".ssh")
}

func getDefaultKeyPath() string {
	return filepath.Join(getDefaultSSHDir(), "id_rsa")
}

func getDefaultPubKeyPath() string {
	return filepath.Join(getDefaultSSHDir(), "id_rsa.pub")
}

// readPassword - hides input when typing password
// fallback to regular input if terminal detection fails (like in pipes)
func readPassword(prompt string) (string, error) {
	fmt.Print(prompt)

	if term.IsTerminal(int(os.Stdin.Fd())) {
		password, err := term.ReadPassword(int(os.Stdin.Fd()))
		fmt.Println()
		return string(password), err
	}

	// not a terminal, just read normally (useful for scripting i guess)
	reader := bufio.NewReader(os.Stdin)
	password, err := reader.ReadString('\n')
	return strings.TrimSpace(password), err
}

// injectSSHKey - the main thing this tool does
// connects to remote server with password and adds your pubkey to authorized_keys
func injectSSHKey(reader *bufio.Reader) {
	fmt.Println(colorCyan + "\n  ═══════════════════════════════════════════════════════" + colorReset)
	fmt.Println(colorWhite + "               SSH KEY INJECTION" + colorReset)
	fmt.Println(colorCyan + "  ═══════════════════════════════════════════════════════" + colorReset)

	// get all the info we need from user
	fmt.Print(colorYellow + "\n  Enter remote server IP/hostname: " + colorReset)
	host, _ := reader.ReadString('\n')
	host = strings.TrimSpace(host)

	if host == "" {
		fmt.Println(colorRed + "  Error: IP/hostname cannot be empty" + colorReset)
		return
	}

	fmt.Print(colorYellow + "  Enter SSH port [22]: " + colorReset)
	portStr, _ := reader.ReadString('\n')
	portStr = strings.TrimSpace(portStr)
	if portStr == "" {
		portStr = "22"
	}

	fmt.Print(colorYellow + "  Enter username: " + colorReset)
	username, _ := reader.ReadString('\n')
	username = strings.TrimSpace(username)

	if username == "" {
		fmt.Println(colorRed + "  Error: Username cannot be empty" + colorReset)
		return
	}

	password, err := readPassword(colorYellow + "  Enter password: " + colorReset)
	if err != nil {
		fmt.Println(colorRed + "  Error reading password: " + err.Error() + colorReset)
		return
	}

	// need to know what os the remote is running
	// the commands are different for windows vs unix
	fmt.Println(osSelectArt)
	fmt.Print(colorYellow + "  Select remote OS type: " + colorReset)
	osChoice, _ := reader.ReadString('\n')
	osChoice = strings.TrimSpace(osChoice)

	var remoteOS string
	switch osChoice {
	case "1":
		remoteOS = "linux"
	case "2":
		remoteOS = "darwin"
	case "3":
		remoteOS = "windows"
	default:
		fmt.Println(colorRed + "  Invalid OS selection" + colorReset)
		return
	}

	// read the public key file
	pubKeyPath := getDefaultPubKeyPath()
	fmt.Print(colorYellow + "  Enter path to public key [" + pubKeyPath + "]: " + colorReset)
	customPath, _ := reader.ReadString('\n')
	customPath = strings.TrimSpace(customPath)
	if customPath != "" {
		pubKeyPath = customPath
	}

	pubKey, err := os.ReadFile(pubKeyPath)
	if err != nil {
		fmt.Println(colorRed + "  Error reading public key: " + err.Error() + colorReset)
		fmt.Println(colorYellow + "  Tip: Generate a key pair first using option [2]" + colorReset)
		return
	}

	pubKeyStr := strings.TrimSpace(string(pubKey))

	fmt.Println(colorCyan + "\n  Connecting to " + host + ":" + portStr + "..." + colorReset)

	// set up ssh connection with password auth
	config := &ssh.ClientConfig{
		User: username,
		Auth: []ssh.AuthMethod{
			ssh.Password(password),
		},
		// TODO: maybe add proper host key checking later
		// for now just accept anything, its fine for most use cases
		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
		Timeout:         30 * time.Second,
	}

	client, err := ssh.Dial("tcp", net.JoinHostPort(host, portStr), config)
	if err != nil {
		fmt.Println(colorRed + "  Connection failed: " + err.Error() + colorReset)
		return
	}
	defer client.Close()

	fmt.Println(colorGreen + "  Connected successfully!" + colorReset)

	// build the command to inject the key
	// different approach for windows vs linux/mac
	var cmd string
	switch remoteOS {
	case "linux", "darwin":
		// unix is easy - just append to authorized_keys
		// grep first to avoid duplicates
		cmd = fmt.Sprintf(`
mkdir -p ~/.ssh && \
chmod 700 ~/.ssh && \
touch ~/.ssh/authorized_keys && \
chmod 600 ~/.ssh/authorized_keys && \
grep -qxF '%s' ~/.ssh/authorized_keys 2>/dev/null || echo '%s' >> ~/.ssh/authorized_keys
`, pubKeyStr, pubKeyStr)

	case "windows":
		// windows openssh is weird
		// admins use C:\ProgramData\ssh\administrators_authorized_keys
		// regular users use %USERPROFILE%\.ssh\authorized_keys
		// just gonna do the user one for now
		cmd = fmt.Sprintf(`
if not exist "%%USERPROFILE%%\.ssh" mkdir "%%USERPROFILE%%\.ssh"
findstr /x /c:"%s" "%%USERPROFILE%%\.ssh\authorized_keys" >nul 2>&1 || echo %s >> "%%USERPROFILE%%\.ssh\authorized_keys"
`, pubKeyStr, pubKeyStr)
	}

	session, err := client.NewSession()
	if err != nil {
		fmt.Println(colorRed + "  Failed to create session: " + err.Error() + colorReset)
		return
	}
	defer session.Close()

	fmt.Println(colorCyan + "  Injecting SSH key..." + colorReset)

	// run the command
	var output []byte
	if remoteOS == "windows" {
		output, err = session.CombinedOutput("cmd /c " + cmd)
	} else {
		output, err = session.CombinedOutput(cmd)
	}

	if err != nil {
		fmt.Println(colorRed + "  Key injection failed: " + err.Error() + colorReset)
		if len(output) > 0 {
			fmt.Println(colorRed + "  Output: " + string(output) + colorReset)
		}
		return
	}

	fmt.Println(colorGreen + "\n  ✓ SSH key injected successfully!" + colorReset)
	fmt.Println(colorGreen + "  You can now login without password:" + colorReset)
	fmt.Println(colorWhite + fmt.Sprintf("    ssh %s@%s", username, host) + colorReset)
}

// generateSSHKey - wrapper around ssh-keygen
// could implement this in pure go but ssh-keygen is everywhere anyway
func generateSSHKey(reader *bufio.Reader) {
	fmt.Println(colorCyan + "\n  ═══════════════════════════════════════════════════════" + colorReset)
	fmt.Println(colorWhite + "             GENERATE SSH KEY PAIR" + colorReset)
	fmt.Println(colorCyan + "  ═══════════════════════════════════════════════════════" + colorReset)

	keyPath := getDefaultKeyPath()
	fmt.Print(colorYellow + "\n  Enter key path [" + keyPath + "]: " + colorReset)
	customPath, _ := reader.ReadString('\n')
	customPath = strings.TrimSpace(customPath)
	if customPath != "" {
		keyPath = customPath
	}

	// dont want to accidentally overwrite someones existing key
	if _, err := os.Stat(keyPath); err == nil {
		fmt.Print(colorYellow + "  Key already exists. Overwrite? (y/N): " + colorReset)
		confirm, _ := reader.ReadString('\n')
		confirm = strings.TrimSpace(strings.ToLower(confirm))
		if confirm != "y" && confirm != "yes" {
			fmt.Println(colorYellow + "  Key generation cancelled." + colorReset)
			return
		}
	}

	// make sure .ssh dir exists
	sshDir := filepath.Dir(keyPath)
	if err := os.MkdirAll(sshDir, 0700); err != nil {
		fmt.Println(colorRed + "  Failed to create .ssh directory: " + err.Error() + colorReset)
		return
	}

	fmt.Print(colorYellow + "  Enter passphrase (empty for none): " + colorReset)
	passphrase, _ := readPassword("")

	fmt.Println(colorCyan + "\n  Generating RSA 4096-bit key pair..." + colorReset)

	// just shell out to ssh-keygen, its the most reliable way
	args := []string{"-t", "rsa", "-b", "4096", "-f", keyPath, "-N", passphrase}
	cmd := exec.Command("ssh-keygen", args...)

	output, err := cmd.CombinedOutput()
	if err != nil {
		fmt.Println(colorRed + "  Key generation failed: " + err.Error() + colorReset)
		fmt.Println(colorYellow + "  Note: ssh-keygen must be available on your system" + colorReset)
		if len(output) > 0 {
			fmt.Println(colorRed + "  Output: " + string(output) + colorReset)
		}
		return
	}

	// fix permissions on unix - ssh is picky about this
	if runtime.GOOS != "windows" {
		os.Chmod(keyPath, 0600)
		os.Chmod(keyPath+".pub", 0644)
	}

	fmt.Println(colorGreen + "\n  ✓ SSH key pair generated successfully!" + colorReset)
	fmt.Println(colorWhite + "    Private key: " + keyPath + colorReset)
	fmt.Println(colorWhite + "    Public key:  " + keyPath + ".pub" + colorReset)
}

// viewPublicKey - just prints your public key so you can copy it
func viewPublicKey() {
	fmt.Println(colorCyan + "\n  ═══════════════════════════════════════════════════════" + colorReset)
	fmt.Println(colorWhite + "              YOUR PUBLIC KEY" + colorReset)
	fmt.Println(colorCyan + "  ═══════════════════════════════════════════════════════" + colorReset)

	pubKeyPath := getDefaultPubKeyPath()

	// check for ed25519 too since thats becoming more common
	if _, err := os.Stat(pubKeyPath); os.IsNotExist(err) {
		ed25519Path := filepath.Join(getDefaultSSHDir(), "id_ed25519.pub")
		if _, err := os.Stat(ed25519Path); err == nil {
			pubKeyPath = ed25519Path
		}
	}

	pubKey, err := os.ReadFile(pubKeyPath)
	if err != nil {
		fmt.Println(colorRed + "\n  No public key found at: " + pubKeyPath + colorReset)
		fmt.Println(colorYellow + "  Generate a key pair first using option [2]" + colorReset)
		return
	}

	fmt.Println(colorGreen + "\n  Public key (" + pubKeyPath + "):" + colorReset)
	fmt.Println(colorCyan + "  ───────────────────────────────────────────────────────" + colorReset)
	fmt.Println(colorWhite + "  " + strings.TrimSpace(string(pubKey)) + colorReset)
	fmt.Println(colorCyan + "  ───────────────────────────────────────────────────────" + colorReset)
}

// testConnection - verify that key auth is working after injection
func testConnection(reader *bufio.Reader) {
	fmt.Println(colorCyan + "\n  ═══════════════════════════════════════════════════════" + colorReset)
	fmt.Println(colorWhite + "             TEST SSH CONNECTION" + colorReset)
	fmt.Println(colorCyan + "  ═══════════════════════════════════════════════════════" + colorReset)

	fmt.Print(colorYellow + "\n  Enter remote server IP/hostname: " + colorReset)
	host, _ := reader.ReadString('\n')
	host = strings.TrimSpace(host)

	if host == "" {
		fmt.Println(colorRed + "  Error: IP/hostname cannot be empty" + colorReset)
		return
	}

	fmt.Print(colorYellow + "  Enter SSH port [22]: " + colorReset)
	portStr, _ := reader.ReadString('\n')
	portStr = strings.TrimSpace(portStr)
	if portStr == "" {
		portStr = "22"
	}

	fmt.Print(colorYellow + "  Enter username: " + colorReset)
	username, _ := reader.ReadString('\n')
	username = strings.TrimSpace(username)

	if username == "" {
		fmt.Println(colorRed + "  Error: Username cannot be empty" + colorReset)
		return
	}

	keyPath := getDefaultKeyPath()
	fmt.Print(colorYellow + "  Enter path to private key [" + keyPath + "]: " + colorReset)
	customPath, _ := reader.ReadString('\n')
	customPath = strings.TrimSpace(customPath)
	if customPath != "" {
		keyPath = customPath
	}

	keyData, err := os.ReadFile(keyPath)
	if err != nil {
		fmt.Println(colorRed + "  Error reading private key: " + err.Error() + colorReset)
		return
	}

	// try parsing the key, might need passphrase
	signer, err := ssh.ParsePrivateKey(keyData)
	if err != nil {
		// probably encrypted, ask for passphrase
		fmt.Print(colorYellow + "  Enter key passphrase: " + colorReset)
		passphrase, _ := readPassword("")
		signer, err = ssh.ParsePrivateKeyWithPassphrase(keyData, []byte(passphrase))
		if err != nil {
			fmt.Println(colorRed + "  Failed to parse private key: " + err.Error() + colorReset)
			return
		}
	}

	fmt.Println(colorCyan + "\n  Testing connection to " + host + ":" + portStr + "..." + colorReset)

	config := &ssh.ClientConfig{
		User: username,
		Auth: []ssh.AuthMethod{
			ssh.PublicKeys(signer),
		},
		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
		Timeout:         30 * time.Second,
	}

	client, err := ssh.Dial("tcp", net.JoinHostPort(host, portStr), config)
	if err != nil {
		fmt.Println(colorRed + "  Connection failed: " + err.Error() + colorReset)
		fmt.Println(colorYellow + "  Make sure the public key is in the remote authorized_keys" + colorReset)
		return
	}
	defer client.Close()

	// run a quick command to make sure everything works
	session, err := client.NewSession()
	if err != nil {
		fmt.Println(colorRed + "  Failed to create session: " + err.Error() + colorReset)
		return
	}
	defer session.Close()

	output, err := session.CombinedOutput("echo 'SSH key authentication successful'")
	if err != nil {
		fmt.Println(colorRed + "  Command failed: " + err.Error() + colorReset)
		return
	}

	fmt.Println(colorGreen + "\n  ✓ " + strings.TrimSpace(string(output)) + colorReset)
	fmt.Println(colorGreen + "  Passwordless SSH login is working!" + colorReset)
}

// interactiveShell - might add this to menu later
// opens a full interactive shell session
func interactiveShell(client *ssh.Client) error {
	session, err := client.NewSession()
	if err != nil {
		return err
	}
	defer session.Close()

	session.Stdout = os.Stdout
	session.Stderr = os.Stderr
	session.Stdin = os.Stdin

	modes := ssh.TerminalModes{
		ssh.ECHO:          1,
		ssh.TTY_OP_ISPEED: 14400,
		ssh.TTY_OP_OSPEED: 14400,
	}

	if err := session.RequestPty("xterm-256color", 80, 40, modes); err != nil {
		return err
	}

	if err := session.Shell(); err != nil {
		return err
	}

	return session.Wait()
}

// unused for now but keeping it around
func copyConn(writer io.Writer, reader io.Reader) {
	io.Copy(writer, reader)
}