// ABOUTME: CLI tool that mounts remote filesystems via sshfs based on
// ABOUTME: ~/.ssh/config entries, with optional editor launch on the mountpoint.

package main

import (
	"flag"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"

	"github.com/kevinburke/ssh_config"
	"github.com/moby/sys/mountinfo"
)

var (
	rxHostUser     = regexp.MustCompile(`^[A-Za-z0-9_][A-Za-z0-9._-]*$`)
	rxIdentityFile = regexp.MustCompile(`^[A-Za-z0-9/~][A-Za-z0-9._@/:+=~%-]*$`)
	rxRemoteDir    = regexp.MustCompile(`^[A-Za-z0-9/~][A-Za-z0-9._/~-]*$`)
)

var (
	eFlag = flag.Bool("e", false, "open mountpoint in your editor")
	vFlag = flag.Bool("v", false, "verbose: print resolved ssh_config fields")
	rDir  string
)

func init() {
	flag.StringVar(&rDir, "r", "", "remote directory to mount (default: remote home)")
	flag.StringVar(&rDir, "remote-dir", "", "remote directory to mount (default: remote home)")
}

func main() {
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr,
			"usage: sshfsc [flags] <Host>\n\n"+
				"Mount a remote home directory via sshfs based on ~/.ssh/config entries.\n\n"+
				"Flags:\n")
		flag.PrintDefaults()
	}
	flag.Parse()
	args := flag.Args()

	if len(args) == 0 {
		fmt.Fprintln(os.Stderr, "No hostname specified.")
		os.Exit(2)
	}

	alias := args[0]
	if err := validate_ssh_field("alias", alias, rxHostUser); err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(3)
	}

	hostname, err := lookup_ssh_field(args[0], "HostName")
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(3)
	}
	user, err := lookup_ssh_field(args[0], "User")
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(3)
	}
	port, err := lookup_ssh_field(args[0], "Port")
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(3)
	}
	ifile, err := lookup_ssh_field(args[0], "IdentityFile")
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(3)
	}

	if hostname == "" || user == "" || ifile == "" {
		var missing []string
		if hostname == "" {
			missing = append(missing, "HostName")
		}
		if user == "" {
			missing = append(missing, "User")
		}
		if ifile == "" {
			missing = append(missing, "IdentityFile")
		}
		fmt.Fprintf(os.Stderr, "ssh config %q missing required field(s): %s\n",
			args[0], strings.Join(missing, ", "))
		os.Exit(3)
	}

	if err := validate_ssh_field("HostName", hostname, rxHostUser); err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(3)
	}
	if err := validate_ssh_field("User", user, rxHostUser); err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(3)
	}
	if err := validate_ssh_field("IdentityFile", ifile, rxIdentityFile); err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(3)
	}
	if p, perr := strconv.Atoi(port); perr != nil || p < 1 || p > 65535 {
		fmt.Fprintf(os.Stderr, "ssh config Port %q is not a valid port number\n", port)
		os.Exit(3)
	}
	if rDir != "" {
		if err := validate_ssh_field("remote-dir", rDir, rxRemoteDir); err != nil {
			fmt.Fprintln(os.Stderr, err)
			os.Exit(3)
		}
	}

	mount, err := verify_mount_dir(alias)
	if err != nil {
		fmt.Fprintln(os.Stderr, "verify_mount_dir() failed:", err)
		os.Exit(4)
	}

	if *vFlag {
		fmt.Println("Hostname:	", hostname)
		fmt.Println("User:		", user)
		fmt.Println("Port:		", port)
		fmt.Println("Ifile:		", ifile)
		fmt.Println("Mount:		", mount)
		if rDir != "" {
			fmt.Println("Remote:		", rDir)
		}
		fmt.Println("---")
	} else {
		fmt.Println("Mount:	", mount)
	}

	chkmount, chkmount_err := mountinfo.Mounted(mount)
	if chkmount_err != nil {
		fmt.Fprintf(os.Stderr, "mountinfo.Mounted() failed with %s\n", chkmount_err)
		os.Exit(5)
	}
	if !chkmount {
		if err := mount_sshfs(alias, user, ifile, port, mount, rDir); err != nil {
			fmt.Fprintln(os.Stderr, "mount_sshfs() failed:", err)
			os.Exit(6)
		}
	} else {
		fmt.Println("!!! Already mounted")
	}
	// run_editor fires on -e regardless of mount state: re-invoking with -e
	// re-opens the editor on an already-mounted server.
	if err := run_editor(mount); err != nil {
		fmt.Fprintln(os.Stderr, "run_editor() failed:", err)
		os.Exit(7)
	}
}

func lookup_ssh_field(alias, field string) (string, error) {
	v, err := ssh_config.GetStrict(alias, field)
	if err != nil {
		return "", fmt.Errorf("parse ~/.ssh/config: %w", err)
	}
	return v, nil
}

func validate_ssh_field(name, value string, allow *regexp.Regexp) error {
	if !allow.MatchString(value) {
		return fmt.Errorf("ssh config %s %q contains disallowed characters", name, value)
	}
	return nil
}

func run_editor(mount string) error {
	if !*eFlag {
		return nil
	}
	cmd := exec.Command("subl", mount)
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	return cmd.Start()
}

func verify_mount_dir(name string) (string, error) {
	runtime := os.Getenv("XDG_RUNTIME_DIR")
	if runtime == "" {
		runtime = filepath.Join("/run/user", strconv.Itoa(os.Getuid()))
	}
	base := filepath.Join(runtime, "sshfs")
	if err := os.MkdirAll(base, 0700); err != nil {
		return "", fmt.Errorf("create base %q: %w", base, err)
	}
	realBase, err := filepath.EvalSymlinks(base)
	if err != nil {
		return "", fmt.Errorf("resolve base %q: %w", base, err)
	}
	mount := filepath.Clean(filepath.Join(realBase, name))
	if !strings.HasPrefix(mount, realBase+string(os.PathSeparator)) {
		return "", fmt.Errorf("name %q escapes mount base %q", name, realBase)
	}
	if info, err := os.Lstat(mount); err == nil && info.Mode()&os.ModeSymlink != 0 {
		return "", fmt.Errorf("mount path %q is a symlink", mount)
	}
	if err := os.MkdirAll(mount, 0700); err != nil {
		return "", fmt.Errorf("create mount dir %q: %w", mount, err)
	}
	return mount, nil
}

func mount_sshfs(alias string, user string, ifile string, port string, mount string, remoteDir string) error {
	cmd := exec.Command("sshfs", "-p", port,
		"-o", "IdentityFile="+ifile,
		"-o", "idmap=user",
		"-o", "cache=yes",
		"-o", "auto_cache",
		"-o", "attr_timeout=60",
		"-o", "entry_timeout=60",
		"-o", "negative_timeout=20",
		"-o", "Ciphers=aes128-gcm@openssh.com",
		"-o", "Compression=no",
		"-o", "reconnect",
		"-o", "StrictHostKeyChecking=accept-new",
		"-o", "ServerAliveInterval=15",
		"-o", "ServerAliveCountMax=3",
		user+"@"+alias+":"+remoteDir, mount)
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	return cmd.Run()
}