From b080b7582bbceee5d9ab5ad3fb0dd784dc276a51 Mon Sep 17 00:00:00 2001
From: max furman <mx.furman@gmail.com>
Date: Tue, 9 Nov 2021 15:26:34 -0800
Subject: [PATCH] Template updates to support multiple SSH include snippets

---
 templates/templates.go | 20 ++++++++++++++------
 templates/values.go    | 11 +++++++----
 2 files changed, 21 insertions(+), 10 deletions(-)

diff --git a/templates/templates.go b/templates/templates.go
index 4fb831d4..aa77a2ff 100644
--- a/templates/templates.go
+++ b/templates/templates.go
@@ -20,6 +20,8 @@ type TemplateType string
 const (
 	// Snippet will mark a template as a part of a file.
 	Snippet TemplateType = "snippet"
+	// FullSnippet will mark a template that includes header and footer as a part of a file.
+	FullSnippet TemplateType = "fullSnippet"
 	// File will mark a templates as a full file.
 	File TemplateType = "file"
 	// Directory will mark a template as a directory.
@@ -99,7 +101,7 @@ func (t *SSHTemplates) Validate() (err error) {
 	return
 }
 
-// Template represents on template file.
+// Template represents a template file.
 type Template struct {
 	*template.Template
 	Name         string       `json:"name"`
@@ -118,8 +120,8 @@ func (t *Template) Validate() error {
 		return nil
 	case t.Name == "":
 		return errors.New("template name cannot be empty")
-	case t.Type != Snippet && t.Type != File && t.Type != Directory:
-		return errors.Errorf("invalid template type %s, it must be %s, %s, or %s", t.Type, Snippet, File, Directory)
+	case t.Type != Snippet && t.Type != File && t.Type != Directory && t.Type != FullSnippet:
+		return errors.Errorf("invalid template type %s, it must be %s, %s, %s, or %s", t.Type, Snippet, FullSnippet, File, Directory)
 	case t.TemplatePath == "" && t.Type != Directory && len(t.Content) == 0:
 		return errors.New("template template cannot be empty")
 	case t.TemplatePath != "" && t.Type == Directory:
@@ -257,11 +259,17 @@ func (o *Output) Write() error {
 		return err
 	}
 
-	if o.Type == File {
+	switch o.Type {
+	case File:
 		return fileutil.WriteFile(path, o.Content, 0600)
+	case Snippet:
+		return fileutil.WriteSnippet(path, o.Content, 0600)
+	case FullSnippet:
+		lines := strings.Split(string(o.Content), "\n")
+		return fileutil.WriteFullSnippet(path, o.Content, lines[0], lines[len(lines)-1], 0600)
+	default:
+		return errors.Errorf("unexpected output template type %s", string(o.Type))
 	}
-
-	return fileutil.WriteSnippet(path, o.Content, 0600)
 }
 
 func mkdir(path string, perm os.FileMode) error {
diff --git a/templates/values.go b/templates/values.go
index 972b1d55..0117cf52 100644
--- a/templates/values.go
+++ b/templates/values.go
@@ -23,7 +23,7 @@ var DefaultSSHTemplates = SSHTemplates{
 	User: []Template{
 		{
 			Name:         "include.tpl",
-			Type:         Snippet,
+			Type:         FullSnippet,
 			TemplatePath: "templates/ssh/include.tpl",
 			Path:         "~/.ssh/config",
 			Comment:      "#",
@@ -67,18 +67,21 @@ var DefaultSSHTemplateData = map[string]string{
 	// include.tpl adds the step ssh config file.
 	//
 	// Note: on windows `Include C:\...` is treated as a relative path.
-	"include.tpl": `Host *
+	"include.tpl": `{{- if .User.Authority }}# {{ .User.Authority }}
+{{ end }}# autogenerated by step
+# @ {{ now }}
+Host *
 {{- if or .User.GOOS "none" | eq "windows" }}
 	Include "{{ .User.StepPath | replace "\\" "/" | trimPrefix "C:" }}/ssh/config"
 {{- else }}
 	Include "{{.User.StepPath}}/ssh/config"
-{{- end }}`,
+{{ end }}# end`,
 
 	// config.tpl is the step ssh config file, it includes the Match rule and
 	// references the step known_hosts file.
 	//
 	// Note: on windows ProxyCommand requires the full path
-	"config.tpl": `Match exec "step ssh check-host %h"
+	"config.tpl": `Match exec "step ssh{{- if .User.Context }} --context {{ .User.Context }}{{- end }} check-host %h"
 {{- if .User.User }}
 	User {{.User.User}}
 {{- end }}