work/app.go

package main

import (
	"fmt"
	"log/slog"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"time"

	"github.com/charmbracelet/huh"
	"golang.org/x/crypto/ssh"
	"golang.org/x/crypto/ssh/knownhosts"
)

type App struct {
	cfg       Config
	flags     Flags
	timeStore *TimeStore
}

func NewApp() (*App, error) {
	cfg, err := loadConfig()
	if err != nil {
		return nil, fmt.Errorf("error loading config: %w", err)
	}

	ts, err := NewTimeStore(cfg)
	if err != nil {
		return nil, fmt.Errorf("error initializing time store: %w", err)
	}

	return &App{
		cfg:       cfg,
		timeStore: ts,
	}, nil
}

func (a *App) Close() error {
	if a.timeStore != nil {
		return a.timeStore.Close()
	}
	return nil
}

func (a *App) StartTracking(tag string, withoutTimew bool) error {
	if err := a.timeStore.StartTracking(tag); err != nil {
		return err
	}

	if !withoutTimew {
		if err := a.runCommand("timew", "start", tag); err != nil {
			slog.Warn("Failed to start timewarrior (ignoring)", "error", err)
		}
	}
	return nil
}

func (a *App) StopTracking(withoutTimew bool) error {
	if err := a.timeStore.StopTracking(); err != nil {
		return err
	}

	if !withoutTimew {
		if err := a.runCommand("timew", "stop"); err != nil {
			slog.Warn("Failed to stop timewarrior (ignoring)", "error", err)
		}
	}
	return nil
}

func (a *App) connect(withoutTimew bool) (*SSHConnection, error) {
	if err := a.StartTracking(TagWork, withoutTimew); err != nil {
		slog.Warn(fmt.Sprintf("Failed to start time tracking for '%s': %v", TagWork, err))
	}

	a.wakeWorkstation()

	sshCon, err := a.newSSHConnection()
	if err != nil {
		return nil, fmt.Errorf("failed to establish primary SSH connection: %w", err)
	}

	slog.Info("SSH connection established. Setting up tunnels...")

	sshForwarder := NewPortForwarder(sshCon.client, "2048", "22", a.cfg.WorkstationIP)
	go func() {
		slog.Info("Starting SSH forwarder (local :2048 -> remote workstation:22)")
		if err := sshForwarder.forward(); err != nil {
			slog.Error(fmt.Sprintf("SSH forwarder failed: %v", err))
		}
		slog.Info("SSH forwarder stopped.")
	}()

	rdpForwarder := NewPortForwarder(sshCon.client, "6000", "3389", a.cfg.WorkstationIP)
	go func() {
		slog.Info("Starting RDP forwarder (local :6000 -> remote workstation:3389)")
		if err := rdpForwarder.forward(); err != nil {
			slog.Error(fmt.Sprintf("ERROR: RDP forwarder failed: %v", err))
		}
		slog.Info("RDP forwarder stopped.")
	}()

	time.Sleep(500 * time.Millisecond)

	return sshCon, nil
}

func (a *App) runCommand(name string, args ...string) error {
	slog.Info(fmt.Sprintf("Executing command: %s %s", name, strings.Join(args, " ")))
	cmd := exec.Command(name, args...)
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	cmd.Stdin = os.Stdin
	err := cmd.Run()
	if err != nil {
		slog.Error(fmt.Sprintf("Command failed: %s %s -> %v", name, strings.Join(args, " "), err))
		return fmt.Errorf("command execution failed: %w", err)
	}
	slog.Info(fmt.Sprintf("Command finished successfully: %s", name))
	return nil
}

func (a *App) wakeWorkstation() {
	slog.Info("Attempting to wake workstation...")
	innerSSHCmd := fmt.Sprintf("ssh -tt %s@%s \"wakeonlan %s && echo 'Wake-on-LAN packet sent.' && exit\"",
		a.cfg.JumpUser,
		a.cfg.JumpHost,
		a.cfg.WorkstationMac)

	outerSSHCmd := []string{
		"-tt",
		"-p", fmt.Sprintf("%d", a.cfg.SSHPort),
		fmt.Sprintf("%s@%s", a.cfg.SSHUser, a.cfg.SSHHost),
		innerSSHCmd,
	}

	if err := a.runCommand("ssh", outerSSHCmd...); err != nil {
		slog.Warn("Failed to send Wake-on-LAN packet via SSH jump. Workstation might already be awake or command failed.")
	} else {
		slog.Info("Wake-on-LAN command executed.")
	}
}

func (a *App) connectToJump() {
	slog.Info("Connecting to Jump Host with Port Forwarding...")
	sshArgs := []string{
		"-tt",
		"-L", fmt.Sprintf("2048:%s:22", a.cfg.WorkstationHost),
		"-p", fmt.Sprintf("%d", a.cfg.SSHPort),
		fmt.Sprintf("%s@%s", a.cfg.SSHUser, a.cfg.SSHHost),
	}
	if err := a.runCommand("ssh", sshArgs...); err != nil {
	}
}

func (a *App) connectToWorkstation() {
	slog.Info("Connecting to Workstation via local tunnel (localhost:2048)...")
	sshArgs := []string{
		"-tt",
		"-L", fmt.Sprintf("6000:%s:3389", a.cfg.WorkstationHost),
		"-p", "2048",
		fmt.Sprintf("%s@127.0.0.1", a.cfg.WorkstationUser),
	}
	if err := a.runCommand("ssh", sshArgs...); err != nil {
		return
	}
}

func (a *App) startRDPConnection() {
	slog.Info("Starting RDP connection to localhost:6000...")

	args := []string{
		fmt.Sprintf("/u:%s", a.cfg.RDPUser),
		fmt.Sprintf("/p:%s", a.cfg.SSHPassword),
		"/v:127.0.0.1:6000",
		"/size:3000x1350",
		"+clipboard",
		"/dynamic-resolution",
	}

	if err := a.runCommand("xfreerdp", args...); err != nil {
		return
	}
}

func (a *App) makeChoice() {
	var choice string

	form := huh.NewForm(
		huh.NewGroup(
			huh.NewSelect[string]().
				Title("What would you like to do?").
				Options(
					huh.NewOption("Start Work & Connect", "start work"),
					huh.NewOption("Stop Work", "stop work"),
					huh.NewOption("Start Break", "start break"),
					huh.NewOption("Stop Break", "stop break"),
					huh.NewOption("Show Today Summary", "show day summary"),
					huh.NewOption("Show Week Summary", "show week summary"),
					huh.NewOption("Show Month Summary", "show month summary"),
					huh.NewOption("Export Yearly Timetable", "export"),
					huh.NewOption("Connect to Jump Host (Tunnel to Workstation)", "connect to jump"),
					huh.NewOption("Connect to Workstation (via Tunnel)", "connect to workstation"),
					huh.NewOption("Start RDP Connection (via Tunnel)", "start rdp connection"),
					huh.NewOption("Wake Workstation", "wake workstation"),
					huh.NewOption("Kill Active Tunnels (Ports 2048, 6000)", "kill tunnels"),
					huh.NewOption("Exit", "exit"),
				).
				Value(&choice),
		),
	)

	err := form.Run()
	if err != nil {
		if err == huh.ErrUserAborted {
			fmt.Println("Operation cancelled.")
			return
		}
		slog.Error(fmt.Sprintf("Form execution failed: %v", err))
		return
	}

	switch choice {
	case "start work":
		a.connect(withoutTimew)
	case "stop work":
		if err := a.StopTracking(withoutTimew); err != nil {
			slog.Error(fmt.Sprintf("Failed to stop time tracking: %v", err))
		}
		if err := a.killForwardings(); err != nil {
			slog.Warn(fmt.Sprintf("Could not kill all forwardings: %v", err))
		}
	case "start break":
		if err := a.StartTracking(TagBreak, withoutTimew); err != nil {
			slog.Error(fmt.Sprintf("Failed to start break tracking: %v", err))
		}
	case "stop break":
		if err := a.StartTracking(TagWork, withoutTimew); err != nil {
			slog.Error(fmt.Sprintf("Failed to stop break (start work): %v", err))
		}
	case "show day summary":
		if err := a.timeStore.ShowSummary("today"); err != nil {
			slog.Error(fmt.Sprintf("Failed to show day summary: %v", err))
		}
	case "show week summary":
		if err := a.timeStore.ShowSummary("week"); err != nil {
			slog.Error(fmt.Sprintf("ERROR: Failed to show week summary: %v", err))
		}
	case "show month summary":
		if err := a.timeStore.ShowSummary("month"); err != nil {
			slog.Error(fmt.Sprintf("Failed to show month summary: %v", err))
		}
	case "export":
		filename := "Arbeitszeiten_" + time.Now().Format("2006") + ".xlsx"
		if a.flags.ExportName != "" && a.flags.ExportName != "Arbeitszeiten.xlsx" {
			filename = a.flags.ExportName
		}
		if err := a.timeStore.ExportSummary(filename); err != nil {
			slog.Error(fmt.Sprintf("Failed to export summary to '%s': %v", filename, err))
		}
	case "connect to jump":
		a.connectToJump()
	case "connect to workstation":
		a.connectToWorkstation()
	case "start rdp connection":
		a.startRDPConnection()
	case "wake workstation":
		a.wakeWorkstation()
	case "kill tunnels":
		if err := a.killForwardings(); err != nil {
			slog.Error(fmt.Sprintf("Failed to kill forwardings: %v", err))
		} else {
			slog.Info("Attempted to kill processes on ports 2048 and 6000.")
		}
	case "exit":
		fmt.Println("Exiting.")
		return
	default:
		slog.Warn(fmt.Sprintf("Unhandled choice '%s'", choice))
	}

	if choice != "exit" && choice != "connect to jump" && choice != "connect to workstation" && choice != "start rdp connection" {
		fmt.Println("\nPress Enter to continue...")
		fmt.Scanln()
		a.makeChoice()
	}
}

func (a *App) getSSHAuth() ssh.AuthMethod {
	keyPath := os.ExpandEnv("$HOME/.ssh/hegenberg")

	keyBytes, err := os.ReadFile(keyPath)
	if err != nil {
		slog.Error(fmt.Sprintf("Unable to read private key '%s': %v", keyPath, err))
		return nil
	}

	var key ssh.Signer
	key, err = ssh.ParsePrivateKey(keyBytes)
	if err != nil {
		if _, ok := err.(*ssh.PassphraseMissingError); ok {
			slog.Info(fmt.Sprintf("Private key '%s' requires a passphrase. Trying with RDP password from config.", keyPath))
			key, err = ssh.ParsePrivateKeyWithPassphrase(keyBytes, []byte(a.cfg.RDPPassword))
			if err != nil {
				slog.Error(fmt.Sprintf("Unable to parse private key '%s' with passphrase: %v", keyPath, err))
				return nil
			}
		} else {
			slog.Error(fmt.Sprintf("Unable to parse private key '%s': %v", keyPath, err))
			return nil
		}
	}

	slog.Info(fmt.Sprintf("Successfully loaded private key '%s'", keyPath))
	return ssh.PublicKeys(key)
}

func (a *App) newSSHConnection() (*SSHConnection, error) {
	authMethod := a.getSSHAuth()
	if authMethod == nil {
		return nil, fmt.Errorf("SSH authentication method could not be obtained")
	}

	homeDir, err := os.UserHomeDir()
	if err != nil {
		return nil, fmt.Errorf("failed to get user home dir: %w", err)
	}
	knownHostsPath := filepath.Join(homeDir, ".ssh", "known_hosts")

	hostKeyCallback, err := knownhosts.New(knownHostsPath)
	if err != nil {
		slog.Warn("Could not load known_hosts file. Please ensure you have connected to the host manually once to populate it.", "path", knownHostsPath)
		return nil, fmt.Errorf("failed to create host key callback (check your known_hosts file): %w", err)
	}

	sshConfig := &ssh.ClientConfig{
		User:            a.cfg.SSHUser,
		Auth:            []ssh.AuthMethod{authMethod},
		HostKeyCallback: hostKeyCallback,
		Timeout:         10 * time.Second,
	}

	target := fmt.Sprintf("%s:%d", a.cfg.SSHHost, a.cfg.SSHPort)
	slog.Info(fmt.Sprintf("Dialing SSH to %s...", target))

	client, err := ssh.Dial("tcp", target, sshConfig)
	if err != nil {
		return nil, fmt.Errorf("SSH dial to %s failed: %w", target, err)
	}
	slog.Info(fmt.Sprintf("SSH connection to %s successful.", target))

	session, err := client.NewSession()
	if err != nil {
		client.Close()
		return nil, fmt.Errorf("failed to create SSH session check: %w", err)
	}
	session.Close()

	return &SSHConnection{
		client: client,
	}, nil
}

func (a *App) killForwardings() error {
	ports := []string{"2048", "6000"}
	killedSomething := false
	var lastErr error

	slog.Info(fmt.Sprintf("Attempting to kill processes listening on ports: %v", strings.Join(ports, ", ")))

	for _, port := range ports {
		cmd := exec.Command("lsof", "-i", "tcp:"+port, "-t")
		output, err := cmd.Output()
		if err != nil {
			if exitErr, ok := err.(*exec.ExitError); ok && exitErr.ExitCode() == 1 {
				slog.Info(fmt.Sprintf("No process found listening on port %s.", port))
			} else {
				slog.Warn(fmt.Sprintf("'lsof' command failed for port %s: %v", port, err))
				lastErr = fmt.Errorf("lsof failed for port %s: %w", port, err)
			}
			continue
		}

		pids := strings.SplitSeq(strings.TrimSpace(string(output)), "\n")
		for pidStr := range pids {
			pid := strings.TrimSpace(pidStr)
			if pid == "" {
				continue
			}
			slog.Info(fmt.Sprintf("Found process PID %s on port %s. Attempting to kill...", pid, port))
			killCmd := exec.Command("kill", pid)
			if err := killCmd.Run(); err != nil {
				slog.Warn(fmt.Sprintf("Failed to kill PID %s (port %s): %v. Trying kill -9...", pid, port, err))
				forceKillCmd := exec.Command("kill", "-9", pid)
				if err := forceKillCmd.Run(); err != nil {
					slog.Error(fmt.Sprintf("Failed to force kill PID %s (port %s): %v", pid, port, err))
					lastErr = fmt.Errorf("kill -9 failed for PID %s: %w", pid, err)
				} else {
					slog.Info(fmt.Sprintf("Force killed PID %s (port %s).", pid, port))
					killedSomething = true
				}
			} else {
				slog.Info(fmt.Sprintf("Killed PID %s (port %s).", pid, port))
				killedSomething = true
			}
		}
	}

	if killedSomething {
		slog.Info("Finished attempting to kill forwarding processes.")
	} else {
		slog.Info("No forwarding processes found or killed.")
	}

	return lastErr
}