cmd/netssh/main.go

@@ -710,9 +710,11 @@ func hookInstallCmd() *cobra.Command {
 	cmd := &cobra.Command{
 		Use:   "install",
 		Short: "Add netssh shell-init to your shell profile",
-		Long: `Appends a single line to your shell profile that runs 'netssh shell-init'
-on every new shell session. shell-init checks whether the cache is older
-than 24 hours and, if so, starts a background refresh — no delay to your prompt.
+		Long: `Adds a marked block (# >>> netssh >>> … # <<< netssh <<<) to your shell
+profile that runs 'netssh shell-init' on every new shell session and, for zsh,
+registers the completion. Unmarked entries from older versions are replaced.
+shell-init checks whether the cache is older than 24 hours and, if so, starts
+a background refresh — no delay to your prompt.
 
 After installation, reload your profile or open a new shell.`,
 		RunE: func(cmd *cobra.Command, args []string) error {
@@ -724,17 +726,12 @@ After installation, reload your profile or open a new shell.`,
 				return err
 			}
 
-			installed, err := hook.Install(profile)
+			// The managed block also sources the aliases (bash/zsh) and registers
+			// zsh completion; fish uses conf.d auto-sourcing instead.
+			installed, err := hook.Install(profile, shell)
 			if err != nil {
 				return err
 			}
-			// For bash/zsh: also add the source line so aliases are loaded on startup.
-			// Fish uses conf.d auto-sourcing — no explicit source line needed.
-			if shell != "fish" {
-				if _, err := hook.InstallAliasesSource(profile); err != nil {
-					return err
-				}
-			}
 
 			if !installed {
 				fmt.Printf("Hook already installed in %s\n", profile)

install.sh

@@ -168,6 +168,11 @@ plan_profile() {
         printf '%s shell-init\n' "$(quote_sh "$INSTALL_DIR/$BINARY")"
         printf '[ ! -f %s ] || . %s\n' "$(quote_sh "$CACHE_DIR/netssh/aliases.sh")" "$(quote_sh "$CACHE_DIR/netssh/aliases.sh")"
       fi
+      if [[ "$kind" == zsh ]]; then
+        # Register completion also when compinit already ran (e.g. oh-my-zsh).
+        printf 'fpath=(%s $fpath)\n' "$(quote_sh "$HOME/.zfunc")"
+        printf 'if (( $+functions[compdef] )); then autoload -Uz _netssh && compdef _netssh netssh; fi\n'
+      fi
       printf '# <<< netssh <<<\n'
     } >> "$data"
   fi
@@ -231,10 +236,25 @@ for ((index=0; index<${#PROFILE_PATHS[@]}; index++)); do
   mv -f "$staged" "$target"
   info "updated  : ${PROFILE_PATHS[$index]}"
 done
+COMPLETIONS=()
 for kind in "${SHELLS[@]}"; do
-  if ! "$INSTALL_DIR/$BINARY" completion install --shell "$kind"; then
+  if "$INSTALL_DIR/$BINARY" completion install --shell "$kind" >/dev/null; then
+    COMPLETIONS+=("$kind")
+  else
     info "warning: could not update $kind completion"
   fi
 done
 printf '\nnetssh %s installed / updated at %s/%s\n' "$TAG" "$INSTALL_DIR" "$BINARY"
+for kind in "${COMPLETIONS[@]}"; do
+  case "$kind" in
+    bash) info "completion: bash -> $HOME/.local/share/bash-completion/completions/netssh" ;;
+    zsh) info "completion: zsh  -> $HOME/.zfunc/_netssh (registered in ${ZSH_DIR}/.zshrc)" ;;
+    fish) info "completion: fish -> $CONFIG_DIR/fish/completions/netssh.fish" ;;
+  esac
+done
 info "Open a new terminal to load the current hooks and aliases."
+for kind in "${COMPLETIONS[@]}"; do
+  if [[ "$kind" == zsh ]]; then
+    info "zsh: if <Tab> still shows files, run: rm -f $(quote_sh "$ZSH_DIR")/.zcompdump* && exec zsh"
+  fi
+done

installer_test.go

@@ -108,7 +108,14 @@ func TestInstallerUpdatesAndConsolidatesProfilesWithoutPython(t *testing.T) {
 			t.Fatalf("bad managed blocks: %s", data)
 		}
 	}
+	zshrc, _ := os.ReadFile(filepath.Join(f.home, ".zshrc"))
+	if !strings.Contains(string(zshrc), "compdef _netssh netssh") || !strings.Contains(string(zshrc), "/.zfunc' $fpath)") {
+		t.Fatalf("zsh completion not registered: %s", zshrc)
+	}
 	data, _ := os.ReadFile(bashrc)
+	if strings.Contains(string(data), "compdef") {
+		t.Fatalf("zsh code in bash profile: %s", data)
+	}
 	for _, kept := range []string{"export KEEP=yes", "# netssh shell-init is documented here", "echo 'netssh shell-init'"} {
 		if !strings.Contains(string(data), kept) {
 			t.Fatalf("unrelated line lost: %s", kept)

internal/hook/hook.go

@@ -62,27 +62,63 @@ func IsInstalled(profilePath string) bool {
 	return false
 }
 
-// Install appends the hook line to the profile if it is not already present.
-// Returns (true, nil) when newly installed, (false, nil) when already present.
+// Install writes the managed netssh block into the profile. Entries written by
+// older versions without markers are replaced instead of duplicated; an
+// existing managed block (e.g. from install.sh) is left untouched.
+// Returns (true, nil) when the profile changed, (false, nil) when already present.
 // Creates the profile file (and any parent directories) if they do not exist.
-func Install(profilePath string) (installed bool, err error) {
-	if IsInstalled(profilePath) {
-		return false, nil
+func Install(profilePath, shell string) (installed bool, err error) {
+	data, err := os.ReadFile(profilePath)
+	if err != nil && !os.IsNotExist(err) {
+		return false, fmt.Errorf("reading profile: %w", err)
 	}
+	for _, line := range strings.Split(string(data), "\n") {
+		if strings.TrimSpace(line) == BlockStart {
+			return false, nil
+		}
+	}
+	cleaned, _, err := clean(string(data), profilePath)
+	if err != nil {
+		return false, err
+	}
+	if cleaned != "" && !strings.HasSuffix(cleaned, "\n") {
+		cleaned += "\n"
+	}
+	if cleaned != "" {
+		cleaned += "\n"
+	}
+	cleaned += Block(shell)
 	if err := os.MkdirAll(filepath.Dir(profilePath), 0o755); err != nil {
 		return false, fmt.Errorf("creating directory: %w", err)
 	}
-	f, err := os.OpenFile(profilePath, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0o644)
-	if err != nil {
-		return false, fmt.Errorf("opening profile: %w", err)
-	}
-	defer f.Close()
-	if _, err := fmt.Fprintf(f, "\n%s\n", Line); err != nil {
+	if err := writeProfile(profilePath, cleaned); err != nil {
 		return false, err
 	}
 	return true, nil
 }
 
+// Block returns the managed profile block for shell.
+func Block(shell string) string {
+	var b strings.Builder
+	b.WriteString(BlockStart + "\n")
+	b.WriteString("netssh shell-init\n")
+	if shell != "fish" {
+		aliases := shellQuote(AliasesPath(), "bash")
+		fmt.Fprintf(&b, "[ ! -f %s ] || . %s\n", aliases, aliases)
+	}
+	if shell == "zsh" {
+		b.WriteString(ZshRegistration(ZshFuncDir()))
+	}
+	b.WriteString(BlockEnd + "\n")
+	return b.String()
+}
+
+// ZshRegistration adds dir to fpath and registers the completion even when
+// compinit already ran earlier in .zshrc (e.g. via oh-my-zsh).
+func ZshRegistration(dir string) string {
+	return fmt.Sprintf("fpath=(%s $fpath)\nif (( $+functions[compdef] )); then autoload -Uz _netssh && compdef _netssh netssh; fi\n", shellQuote(dir, "bash"))
+}
+
 // InstallAliasesSource appends the aliases source line to a bash/zsh profile.
 // This is a no-op for fish (which uses conf.d auto-sourcing instead).
 // Returns (true, nil) when newly added, (false, nil) when already present.
@@ -110,6 +146,26 @@ func InstallAliasesSource(profilePath string) (installed bool, err error) {
 var hookLinePattern = regexp.MustCompile(`^(?:[^\s]*/)?netssh\s+shell-init\s*(?:#.*)?$`)
 var aliasesLinePattern = regexp.MustCompile(`^(?:source|\.)\s+["']?[^;|&]*netssh/aliases\.sh["']?(?:\s+2>/dev/null)?\s*(?:#.*)?$`)
 
+// legacyPatterns match the unmarked lines written by older versions and by
+// manual setup instructions. They mirror the patterns in install.sh.
+var legacyPatterns = []*regexp.Regexp{
+	hookLinePattern,
+	aliasesLinePattern,
+	regexp.MustCompile(`^eval\s+["']?\$\((?:[^\s]*/)?netssh\s+alias(?:\s+--shell\s+(?:bash|zsh))?\)["']?\s*(?:#.*)?$`),
+	regexp.MustCompile(`^(?:[^\s]*/)?netssh\s+alias(?:\s+--shell\s+fish)?\s*\|\s*source\s*(?:#.*)?$`),
+	regexp.MustCompile(`^source\s+<\((?:[^\s]*/)?netssh\s+alias\)\s*(?:#.*)?$`),
+	regexp.MustCompile(`^alias\s+[A-Za-z0-9_.-]+(?:=|\s+)["']netssh(?:\s+[^"']+)?["']\s*(?:#.*)?$`),
+}
+
+func isLegacyLine(line string) bool {
+	for _, pattern := range legacyPatterns {
+		if pattern.MatchString(line) {
+			return true
+		}
+	}
+	return false
+}
+
 // Uninstall removes legacy hook lines and complete installer-managed blocks.
 // Missing profiles are a no-op; malformed blocks leave the file unchanged.
 func Uninstall(profilePath string) (removed bool, err error) {
@@ -120,21 +176,33 @@ func Uninstall(profilePath string) (removed bool, err error) {
 	if err != nil {
 		return false, fmt.Errorf("reading profile: %w", err)
 	}
+	cleaned, removed, err := clean(string(data), profilePath)
+	if err != nil || !removed {
+		return false, err
+	}
+	return true, writeProfile(profilePath, cleaned)
+}
+
+// clean strips managed blocks and legacy lines from profile content.
+func clean(data, profilePath string) (cleaned string, removed bool, err error) {
+	if data == "" {
+		return "", false, nil
+	}
 	var kept []string
 	inside := false
-	for _, line := range strings.Split(string(data), "\n") {
+	for _, line := range strings.Split(data, "\n") {
 		value := strings.TrimSpace(line)
 		switch value {
 		case BlockStart:
 			if inside {
-				return false, fmt.Errorf("nested netssh block in %s", profilePath)
+				return "", false, fmt.Errorf("nested netssh block in %s", profilePath)
 			}
 			inside = true
 			removed = true
 			continue
 		case BlockEnd:
 			if !inside {
-				return false, fmt.Errorf("unmatched netssh block in %s", profilePath)
+				return "", false, fmt.Errorf("unmatched netssh block in %s", profilePath)
 			}
 			inside = false
 			continue
@@ -142,29 +210,41 @@ func Uninstall(profilePath string) (removed bool, err error) {
 		if inside {
 			continue
 		}
-		if hookLinePattern.MatchString(value) || aliasesLinePattern.MatchString(value) {
+		if isLegacyLine(value) {
 			removed = true
 			continue
 		}
 		kept = append(kept, line)
 	}
 	if inside {
-		return false, fmt.Errorf("unclosed netssh block in %s", profilePath)
+		return "", false, fmt.Errorf("unclosed netssh block in %s", profilePath)
 	}
 	if !removed {
-		return false, nil
+		return data, false, nil
 	}
-	cleaned := strings.ReplaceAll(strings.Join(kept, "\n"), "\n\n\n", "\n\n")
-	target, err := filepath.EvalSymlinks(profilePath)
-	if err != nil {
-		return false, err
+	cleaned = strings.ReplaceAll(strings.Join(kept, "\n"), "\n\n\n", "\n\n")
+	// Drop the blank separator line Install puts in front of the block.
+	if strings.HasSuffix(cleaned, "\n\n") {
+		cleaned = strings.TrimRight(cleaned, "\n") + "\n"
 	}
-	info, err := os.Stat(target)
-	if err != nil {
-		return false, err
+	return cleaned, true, nil
+}
+
+// writeProfile replaces the profile atomically, following dotfile symlinks
+// and keeping the permissions of an existing file.
+func writeProfile(profilePath, content string) error {
+	target, err := filepath.EvalSymlinks(profilePath)
+	if os.IsNotExist(err) {
+		target = profilePath
+	} else if err != nil {
+		return err
 	}
-	if err := storage.WriteFile(target, []byte(cleaned), info.Mode().Perm()); err != nil {
-		return false, fmt.Errorf("writing profile: %w", err)
+	mode := os.FileMode(0o644)
+	if info, err := os.Stat(target); err == nil {
+		mode = info.Mode().Perm()
 	}
-	return true, nil
+	if err := storage.WriteFile(target, []byte(content), mode); err != nil {
+		return fmt.Errorf("writing profile: %w", err)
+	}
+	return nil
 }

internal/hook/hook_test.go

@@ -2,6 +2,7 @@ package hook_test
 
 import (
 	"os"
+	"os/exec"
 	"path/filepath"
 	"strings"
 	"testing"
@@ -87,7 +88,7 @@ func TestIsInstalled_MissingFile(t *testing.T) {
 func TestInstall_Fresh(t *testing.T) {
 	profile := filepath.Join(t.TempDir(), ".bashrc")
 
-	installed, err := hook.Install(profile)
+	installed, err := hook.Install(profile, "bash")
 	if err != nil {
 		t.Fatalf("Install: %v", err)
 	}
@@ -102,8 +103,8 @@ func TestInstall_Fresh(t *testing.T) {
 func TestInstall_Idempotent(t *testing.T) {
 	profile := filepath.Join(t.TempDir(), ".bashrc")
 
-	hook.Install(profile)
-	installed, err := hook.Install(profile)
+	hook.Install(profile, "bash")
+	installed, err := hook.Install(profile, "bash")
 	if err != nil {
 		t.Fatalf("second Install: %v", err)
 	}
@@ -120,7 +121,7 @@ func TestInstall_Idempotent(t *testing.T) {
 func TestInstall_CreatesProfileAndParentDirs(t *testing.T) {
 	profile := filepath.Join(t.TempDir(), "nested", "dir", ".zshrc")
 
-	if _, err := hook.Install(profile); err != nil {
+	if _, err := hook.Install(profile, "bash"); err != nil {
 		t.Fatalf("Install should create missing directories: %v", err)
 	}
 	if _, err := os.Stat(profile); err != nil {
@@ -131,7 +132,7 @@ func TestInstall_CreatesProfileAndParentDirs(t *testing.T) {
 func TestInstall_PreservesExistingContent(t *testing.T) {
 	profile := writeProfile(t, "export PATH=$PATH:/usr/local/bin\n")
 
-	hook.Install(profile)
+	hook.Install(profile, "bash")
 
 	data, _ := os.ReadFile(profile)
 	content := string(data)
@@ -145,13 +146,13 @@ func TestInstall_PreservesExistingContent(t *testing.T) {
 
 func TestInstall_AppendedAtEnd(t *testing.T) {
 	profile := writeProfile(t, "existing line\n")
-	hook.Install(profile)
+	hook.Install(profile, "bash")
 
 	data, _ := os.ReadFile(profile)
 	lines := strings.Split(strings.TrimRight(string(data), "\n"), "\n")
 	last := lines[len(lines)-1]
-	if last != hook.Line {
-		t.Errorf("hook line should be the last line, got %q", last)
+	if last != hook.BlockEnd {
+		t.Errorf("managed block should end the profile, got %q", last)
 	}
 }
 
@@ -226,7 +227,7 @@ func TestUninstall_CollapsesExtraBlankLines(t *testing.T) {
 func TestInstallUninstall_Roundtrip(t *testing.T) {
 	profile := writeProfile(t, "existing content\n")
 
-	hook.Install(profile)
+	hook.Install(profile, "bash")
 	if !hook.IsInstalled(profile) {
 		t.Fatal("should be installed after Install")
 	}
@@ -265,7 +266,7 @@ func writeProfile(t *testing.T, content string) string {
 func TestManagedBlockInstallAndUninstall(t *testing.T) {
 	original := "export KEEP=yes\n# netssh shell-init documentation\necho 'netssh shell-init'\n"
 	profile := writeProfile(t, original+hook.BlockStart+"\nexport PATH='/custom/bin':\"$PATH\"\n'/custom/bin/netssh' shell-init\n[ ! -f '/custom/cache/netssh/aliases.sh' ] || . '/custom/cache/netssh/aliases.sh'\n"+hook.BlockEnd+"\n")
-	if installed, err := hook.Install(profile); err != nil || installed {
+	if installed, err := hook.Install(profile, "bash"); err != nil || installed {
 		t.Fatalf("duplicate install: %v, %v", installed, err)
 	}
 	if installed, err := hook.InstallAliasesSource(profile); err != nil || installed {
@@ -288,3 +289,69 @@ func TestProfilePathRespectsZdotdir(t *testing.T) {
 		t.Fatalf("zsh path: %s, %v", path, err)
 	}
 }
+
+func TestInstallReplacesLegacyEntriesWithBlock(t *testing.T) {
+	t.Setenv("XDG_CACHE_HOME", t.TempDir())
+	profile := writeProfile(t, "export KEEP=yes\nnetssh shell-init\nsource ~/.cache/netssh/aliases.sh 2>/dev/null\neval \"$(netssh alias)\"\nalias j='netssh'\necho 'netssh shell-init'\n")
+	installed, err := hook.Install(profile, "zsh")
+	if err != nil || !installed {
+		t.Fatalf("install: %v, %v", installed, err)
+	}
+	data, _ := os.ReadFile(profile)
+	content := string(data)
+	for _, gone := range []string{"\nnetssh shell-init\nsource", "eval ", "alias j="} {
+		if strings.Contains(content, gone) {
+			t.Fatalf("legacy entry %q kept:\n%s", gone, content)
+		}
+	}
+	for _, kept := range []string{"export KEEP=yes", "echo 'netssh shell-init'", hook.BlockStart, "compdef _netssh netssh", "fpath=("} {
+		if !strings.Contains(content, kept) {
+			t.Fatalf("missing %q:\n%s", kept, content)
+		}
+	}
+	if strings.Count(content, hook.BlockStart) != 1 {
+		t.Fatalf("duplicate blocks:\n%s", content)
+	}
+	if removed, err := hook.Uninstall(profile); err != nil || !removed {
+		t.Fatalf("uninstall: %v, %v", removed, err)
+	}
+	data, _ = os.ReadFile(profile)
+	if string(data) != "export KEEP=yes\necho 'netssh shell-init'\n" {
+		t.Fatalf("unexpected remainder: %q", data)
+	}
+}
+
+func TestBlockPerShell(t *testing.T) {
+	if b := hook.Block("fish"); strings.Contains(b, "aliases.sh") || strings.Contains(b, "compdef") {
+		t.Fatalf("fish block: %s", b)
+	}
+	if b := hook.Block("bash"); !strings.Contains(b, "aliases.sh") || strings.Contains(b, "compdef") {
+		t.Fatalf("bash block: %s", b)
+	}
+}
+
+func TestZshBlockRegistersCompletionAfterCompinit(t *testing.T) {
+	zsh, err := exec.LookPath("zsh")
+	if err != nil {
+		t.Skip("zsh unavailable")
+	}
+	home := t.TempDir()
+	t.Setenv("HOME", home)
+	t.Setenv("ZDOTDIR", home)
+	t.Setenv("XDG_CACHE_HOME", filepath.Join(home, "cache"))
+	if err := os.MkdirAll(hook.ZshFuncDir(), 0o755); err != nil {
+		t.Fatal(err)
+	}
+	if err := os.WriteFile(filepath.Join(hook.ZshFuncDir(), "_netssh"), []byte(hook.ZshCompletionScript), 0o644); err != nil {
+		t.Fatal(err)
+	}
+	// compinit runs before the block, as with oh-my-zsh.
+	rc := "autoload -Uz compinit && compinit -u -d \"$HOME/.zcompdump\"\nnetssh() { :; }\n" + hook.Block("zsh")
+	if err := os.WriteFile(filepath.Join(home, ".zshrc"), []byte(rc), 0o644); err != nil {
+		t.Fatal(err)
+	}
+	out, err := exec.Command(zsh, "-i", "-c", `print -r -- $_comps[netssh]`).CombinedOutput()
+	if err != nil || !strings.Contains(string(out), "_netssh") {
+		t.Fatalf("completion not registered: %v\n%s", err, out)
+	}
+}

internal/setup/wizard.go

@@ -193,14 +193,9 @@ func RunWizard(cfg *config.Config) error {
 			fmt.Fprintf(os.Stderr, "Run 'netssh hook install --shell bash|zsh|fish' to activate your alias.\n")
 			return nil
 		}
-		if _, err := hook.Install(profile); err != nil {
+		if _, err := hook.Install(profile, shell); err != nil {
 			return err
 		}
-		if shell != "fish" {
-			if _, err := hook.InstallAliasesSource(profile); err != nil {
-				return err
-			}
-		}
 		fmt.Fprintln(os.Stderr, hook.ReloadNote(profile))
 	}
 	return nil