ssh/config: add method MarshalText and WriteTo

The MarshalText method encode the Section back to ssh_config format
with two spaces as indentation in key.

The WriteTo method marshal the Section into text and write it to
[io.Writer] w.

5 files changed, 156 insertions, 8 deletions

diff --git a/lib/ssh/config/match_criteria.go b/lib/ssh/config/match_criteria.go
index 54de846f..38ecb7b3 100644
--- a/lib/ssh/config/match_criteria.go
+++ b/lib/ssh/config/match_criteria.go
@@ -4,7 +4,11 @@
 
 package config
 
-import "strings"
+import (
+	"bytes"
+	"io"
+	"strings"
+)
 
 const (
 	criteriaAll          = "all"
@@ -47,6 +51,41 @@ func newMatchCriteria(name, arg string) (criteria *matchCriteria, err error) {
 	return criteria, nil
 }
 
+// MarshalText encode the criteria back to ssh_config format.
+func (mcriteria *matchCriteria) MarshalText() (text []byte, err error) {
+	var buf bytes.Buffer
+
+	if mcriteria.isNegate {
+		buf.WriteByte('!')
+	}
+	buf.WriteString(mcriteria.name)
+
+	var (
+		pat *pattern
+		x   int
+	)
+	for x, pat = range mcriteria.patterns {
+		if x == 0 {
+			buf.WriteByte(' ')
+		} else {
+			buf.WriteByte(',')
+		}
+		pat.WriteTo(&buf)
+	}
+
+	return buf.Bytes(), nil
+}
+
+// WriteTo marshal the matchCriteria into text and write it to w.
+func (mcriteria *matchCriteria) WriteTo(w io.Writer) (n int64, err error) {
+	var text []byte
+	text, _ = mcriteria.MarshalText()
+
+	var c int
+	c, err = w.Write(text)
+	return int64(c), err
+}
+
 func (mcriteria *matchCriteria) isMatch(s string) bool {
 	switch mcriteria.name {
 	case criteriaAll:
diff --git a/lib/ssh/config/parser.go b/lib/ssh/config/parser.go
index 8e0d6b03..260f0b46 100644
--- a/lib/ssh/config/parser.go
+++ b/lib/ssh/config/parser.go
@@ -222,10 +222,10 @@ func readLines(file string) (lines []string, err error) {
 }
 
 // parseArgs split single line arguments into list of string, separated by
-// `sep` (default to space), grouped by double quote.
+// sep (default to space), grouped by double quote.
 //
-// For example, given raw argument `a "b c" d` it would return "a", "b c", and
-// "d".
+// For example, given raw argument `a "b c" d` it would return
+// ["a" "b c" "d"].
 func parseArgs(raw string, sep byte) (args []string) {
 	raw = strings.TrimSpace(raw)
 	if len(raw) == 0 {
diff --git a/lib/ssh/config/pattern.go b/lib/ssh/config/pattern.go
index d413fc9a..39199fd8 100644
--- a/lib/ssh/config/pattern.go
+++ b/lib/ssh/config/pattern.go
@@ -4,7 +4,11 @@
 
 package config
 
-import "path/filepath"
+import (
+	"bytes"
+	"io"
+	"path/filepath"
+)
 
 type pattern struct {
 	value    string
@@ -22,6 +26,28 @@ func newPattern(s string) (pat *pattern) {
 	return pat
 }
 
+// MarshalText encode the pattern back to ssh_config format.
+func (pat *pattern) MarshalText() (text []byte, err error) {
+	var buf bytes.Buffer
+
+	if pat.isNegate {
+		buf.WriteByte('!')
+	}
+	buf.WriteString(pat.value)
+
+	return buf.Bytes(), nil
+}
+
+// WriteTo marshal the pattern into text and write it to w.
+func (pat *pattern) WriteTo(w io.Writer) (n int64, err error) {
+	var text []byte
+	text, _ = pat.MarshalText()
+
+	var c int
+	c, err = w.Write(text)
+	return int64(c), err
+}
+
 // isMatch will return true if input string match with regex and isNegate is
 // false; otherwise it will return false.
 func (pat *pattern) isMatch(s string) bool {
diff --git a/lib/ssh/config/section.go b/lib/ssh/config/section.go
index baeeb344..7f6c6126 100644
--- a/lib/ssh/config/section.go
+++ b/lib/ssh/config/section.go
@@ -5,10 +5,13 @@
 package config
 
 import (
+	"bytes"
 	"errors"
 	"fmt"
+	"io"
 	"os"
 	"path/filepath"
+	"sort"
 	"strconv"
 	"strings"
 
@@ -215,6 +218,7 @@ type Section struct {
 	// Criteria for Match section.
 	criteria []*matchCriteria
 
+	// If true indicated that this is Match section.
 	useCriteria bool
 }
 
@@ -570,6 +574,86 @@ func (section *Section) UserKnownHostsFile() []string {
 	return section.knownHostsFile
 }
 
+// MarshalText encode the Section back to ssh_config format.
+// The key is indented by two spaces.
+func (section *Section) MarshalText() (text []byte, err error) {
+	var buf bytes.Buffer
+
+	if section.useCriteria {
+		buf.WriteString(`Match`)
+
+		var criteria *matchCriteria
+		for _, criteria = range section.criteria {
+			buf.WriteByte(' ')
+			criteria.WriteTo(&buf)
+		}
+	} else {
+		buf.WriteString(`Host`)
+
+		if len(section.patterns) == 0 {
+			buf.WriteByte(' ')
+			buf.WriteString(section.name)
+		} else {
+			var pat *pattern
+			for _, pat = range section.patterns {
+				buf.WriteByte(' ')
+				pat.WriteTo(&buf)
+			}
+		}
+	}
+	buf.WriteByte('\n')
+
+	var (
+		listKey = make([]string, 0, len(section.Field))
+		key     string
+		val     string
+	)
+	for key = range section.Field {
+		listKey = append(listKey, key)
+	}
+	sort.Strings(listKey)
+
+	for _, key = range listKey {
+		if key == KeyIdentityFile {
+			for _, val = range section.IdentityFile {
+				buf.WriteString(`  `)
+				buf.WriteString(key)
+				buf.WriteByte(' ')
+				buf.WriteString(section.pathUnfold(val))
+				buf.WriteByte('\n')
+			}
+			continue
+		}
+
+		buf.WriteString(`  `)
+		buf.WriteString(key)
+		buf.WriteByte(' ')
+		buf.WriteString(section.Field[key])
+		buf.WriteByte('\n')
+	}
+
+	return buf.Bytes(), nil
+}
+
+// WriteTo marshal the Section into text and write it to w.
+func (section *Section) WriteTo(w io.Writer) (n int64, err error) {
+	var text []byte
+	text, _ = section.MarshalText()
+
+	var c int
+	c, err = w.Write(text)
+	return int64(c), err
+}
+
+// pathUnfold replace the home directory prefix with '~'.
+func (section *Section) pathUnfold(in string) (out string) {
+	if !strings.HasPrefix(in, section.homeDir) {
+		return in
+	}
+	out = `~` + in[len(section.homeDir):]
+	return out
+}
+
 // setEnv set the Environments with key and value of format "KEY=VALUE".
 func (section *Section) setEnv(env string) {
 	kv := strings.SplitN(env, "=", 2)
diff --git a/lib/ssh/config/section_match.go b/lib/ssh/config/section_match.go
index c9d1eeee..0e30d952 100644
--- a/lib/ssh/config/section_match.go
+++ b/lib/ssh/config/section_match.go
@@ -11,8 +11,7 @@ import (
 )
 
 var (
-	errCriteriaAll = errors.New(`the "all" criteria must appear alone` +
-		` or immediately after "canonical" or "final`)
+	errCriteriaAll = errors.New(`the "all" criteria must appear alone or immediately after "canonical" or "final`)
 )
 
 // newSectionMatch create new Match section using one or more criteria or the
@@ -65,7 +64,7 @@ func newSectionMatch(rawPattern string) (match *Section, err error) {
 			criteria, err = newMatchCriteria(token, arg)
 			x++
 		default:
-			return nil, fmt.Errorf("unknown criteria %q", token)
+			err = fmt.Errorf(`unknown criteria %q`, token)
 		}
 		if err != nil {
 			return nil, err