From a739a16a5e7db9714e274891da2fb4e955f29c4d Mon Sep 17 00:00:00 2001
From: Patryk Hegenberg <patrykhegenberg@gmail.com>
Date: Wed, 2 Apr 2025 22:19:32 +0200
Subject: [PATCH] refactor: perform clean up and add option for direct
 connection

---
 app.go    |  9 ++++----
 cmd.go    | 68 +++++++++++++++++++++++++++++++++++++++++++++++++------
 config.go |  9 ++++----
 3 files changed, 71 insertions(+), 15 deletions(-)

diff --git a/app.go b/app.go
index 7cc7958..33fc50d 100644
--- a/app.go
+++ b/app.go
@@ -42,7 +42,7 @@ func (a *App) Close() error {
 	return nil
 }
 
-func (a *App) connect() {
+func (a *App) connect() (*SSHConnection, error) { // Rückgabetyp geändert
 	if err := a.timeStore.StartTracking(TagWork); err != nil {
 		log.Printf("WARN: Failed to start time tracking for '%s': %v", TagWork, err)
 	}
@@ -51,7 +51,7 @@ func (a *App) connect() {
 
 	sshCon, err := a.newSSHConnection()
 	if err != nil {
-		log.Fatalf("FATAL: Failed to establish primary SSH connection: %v", err)
+		return nil, fmt.Errorf("failed to establish primary SSH connection: %w", err)
 	}
 
 	log.Println("INFO: SSH connection established. Setting up tunnels...")
@@ -74,8 +74,9 @@ func (a *App) connect() {
 		log.Println("INFO: RDP forwarder stopped.")
 	}()
 
-	time.Sleep(1 * time.Second)
-	log.Println("INFO: Tunnels should be active. You can now connect to localhost:2048 (SSH) or localhost:6000 (RDP).")
+	time.Sleep(500 * time.Millisecond)
+
+	return sshCon, nil
 }
 
 func (a *App) runCommand(name string, args ...string) error {
diff --git a/cmd.go b/cmd.go
index 72ef70c..c705acb 100644
--- a/cmd.go
+++ b/cmd.go
@@ -4,7 +4,9 @@ import (
 	"fmt"
 	"log"
 	"os"
+	"os/signal"
 	"strings"
+	"syscall"
 	"time"
 
 	"github.com/spf13/cobra"
@@ -33,17 +35,69 @@ and other utilities.`,
 }
 
 func (a *App) startCommand() *cobra.Command {
-	return &cobra.Command{
+	cmd := &cobra.Command{
 		Use:   "start",
-		Short: "Start work: Track time, wake PC, connect",
-		Long:  "Starts time tracking for 'work', attempts to wake the workstation, and sets up SSH tunnels.",
-		Run: func(cmd *cobra.Command, args []string) {
+		Short: "Start work: Track time, WOL, setup tunnels, optionally connect or run in background",
+		Long: `Starts time tracking, attempts WOL, sets up SSH tunnels.
+Default behavior: Immediately starts an interactive SSH session to the workstation via the tunnel. The command blocks until this session ends.
+Use --background (-b) to keep tunnels running in the background without auto-connecting. Press Ctrl+C to stop background tunnels.`,
+		RunE: func(cmd *cobra.Command, args []string) error {
 			fmt.Println("Starting workday procedures...")
-			a.connect()
-			fmt.Println("Workday start initiated. Tunnels are running in the background.")
-			fmt.Println("Use 'workctl connect rdp' or connect manually.")
+
+			sshCon, err := a.connect()
+			if err != nil {
+				fmt.Printf("ERROR: Failed to start connections: %v\n", err)
+				return fmt.Errorf("connection setup failed: %w", err)
+			}
+
+			defer func() {
+				log.Println("INFO: Closing SSH connection to jump host (defer)...")
+				if err := sshCon.Close(); err != nil {
+					log.Printf("WARN: Error closing SSH connection in defer: %v", err)
+				} else {
+					log.Println("INFO: SSH connection closed via defer.")
+				}
+			}()
+
+			if a.flags.StartInBackground {
+				fmt.Println("\nINFO: Tunnels are active in background.")
+				fmt.Println("       Connect manually via SSH: ssh -p 2048 <user>@127.0.0.1")
+				fmt.Println("       Connect manually via RDP: xfreerdp /v:127.0.0.1:6000 ...")
+				fmt.Println("INFO: Press Ctrl+C to stop tunnels.")
+
+				sigChan := make(chan os.Signal, 1)
+				signal.Notify(sigChan, os.Interrupt, syscall.SIGTERM)
+				<-sigChan
+
+				fmt.Println("\nINFO: Received interrupt signal. Shutting down background process...")
+				log.Println("INFO: Received signal, cleanup via defer sshCon.Close() will run.")
+				if err := a.timeStore.StopTracking(); err != nil {
+					log.Printf("WARN: Failed to stop time tracking: %v", err)
+				} else {
+					log.Println("INFO: Time tracking stopped.")
+				}
+				fmt.Println("INFO: Background shutdown complete.")
+
+			} else {
+				fmt.Println("INFO: Automatically connecting to workstation via SSH tunnel...")
+				a.connectToWorkstation()
+
+				fmt.Println("INFO: Workstation SSH session finished.")
+				log.Println("INFO: Foreground session ended, cleanup via defer sshCon.Close() will run.")
+				if err := a.timeStore.StopTracking(); err != nil {
+					log.Printf("WARN: Failed to stop time tracking: %v", err)
+				} else {
+					log.Println("INFO: Time tracking stopped.")
+				}
+			}
+
+			return nil
 		},
 	}
+
+	cmd.Flags().BoolVarP(&a.flags.StartInBackground, "background", "b", false, "Run tunnels in the background instead of connecting immediately")
+
+	return cmd
 }
 
 func (a *App) stopCommand() *cobra.Command {
diff --git a/config.go b/config.go
index 58da7d5..0800523 100644
--- a/config.go
+++ b/config.go
@@ -26,10 +26,11 @@ type Config struct {
 }
 
 type Flags struct {
-	ShowWeek   bool
-	ShowMonth  bool
-	ShowExport bool
-	ExportName string
+	ShowWeek          bool
+	ShowMonth         bool
+	ShowExport        bool
+	ExportName        string
+	StartInBackground bool
 }
 
 func loadConfig() (Config, error) {