Merge pull request #158 from square/mmc/protocols

Improve FTP & starttls protocol handling

main.go

@@ -49,7 +49,7 @@ var (
 	connectCaPath   = connect.Flag("ca", "Path to CA bundle (system default if unspecified).").ExistingFile()
 	connectCert     = connect.Flag("cert", "Client certificate chain for connecting to server (PEM).").ExistingFile()
 	connectKey      = connect.Flag("key", "Private key for client certificate, if not in same file (PEM).").ExistingFile()
-	connectStartTLS = connect.Flag("start-tls", "Enable StartTLS protocol ('ldap', 'mysql', 'postgres', 'smtp' or 'ftp').").Short('t').PlaceHolder("PROTOCOL").Enum("mysql", "postgres", "psql", "smtp", "ldap", "ftp")
+	connectStartTLS = connect.Flag("start-tls", fmt.Sprintf("Enable StartTLS protocol; one of: %v.", starttls.Protocols)).Short('t').PlaceHolder("PROTOCOL").Enum(starttls.Protocols...)
 	connectIdentity = connect.Flag("identity", "With --start-tls, sets the DB user or SMTP EHLO name").Default("certigo").String()
 	connectProxy    = connect.Flag("proxy", "Optional URI for HTTP(s) CONNECT proxy to dial connections with").URL()
 	connectTimeout  = connect.Flag("timeout", "Timeout for connecting to remote server (can be '5m', '1s', etc).").Default("5s").Duration()
@@ -73,6 +73,9 @@ func main() {
 
 	terminalWidth := determineTerminalWidth()
 
+	// Alias starttls to start-tls
+	connect.Flag("starttls", "").Hidden().EnumVar(connectStartTLS, starttls.Protocols...)
+
 	stdout := colorable.NewColorableStdout()
 	result := simpleResult{}
 	switch kingpin.MustParse(app.Parse(os.Args[1:])) {

starttls/ftp.go

@@ -19,10 +19,9 @@ package starttls
 import (
 	"bufio"
 	"crypto/tls"
-	"errors"
 	"fmt"
 	"net"
-	"strings"
+	"strconv"
 )
 
 func dumpAuthTLSFromFTP(dialer Dialer, address string, config *tls.Config) (*tls.ConnectionState, error) {
@@ -36,8 +35,8 @@ func dumpAuthTLSFromFTP(dialer Dialer, address string, config *tls.Config) (*tls
 	if err != nil {
 		return nil, err
 	}
-	if status != "220" {
-		return nil, fmt.Errorf("FTP server responded with status %s, was expecting 220", status)
+	if status != 220 {
+		return nil, fmt.Errorf("FTP server responded with status %d, was expecting 220", status)
 	}
 
 	fmt.Fprintf(conn, "AUTH TLS\r\n")
@@ -45,8 +44,8 @@ func dumpAuthTLSFromFTP(dialer Dialer, address string, config *tls.Config) (*tls
 	if err != nil {
 		return nil, err
 	}
-	if status != "234" {
-		return nil, fmt.Errorf("FTP server responded with status %s, was expecting 234", status)
+	if status != 234 {
+		return nil, fmt.Errorf("FTP server responded with status %d, was expecting 234", status)
 	}
 
 	tlsConn := tls.Client(conn, config)
@@ -59,15 +58,14 @@ func dumpAuthTLSFromFTP(dialer Dialer, address string, config *tls.Config) (*tls
 	return &state, nil
 }
 
-func readFTP(conn *net.TCPConn) (string, error) {
+func readFTP(conn *net.TCPConn) (int, error) {
 	reader := bufio.NewReader(conn)
 	response, err := reader.ReadString('\n')
 	if err != nil {
-		return "", err
+		return 0, err
 	}
-	status := strings.Split(response, " ")
-	if len(status) == 0 {
-		return "", errors.New("garbled response from FTP server after AUTH TLS command")
+	if len(response) <= 3 {
+		return 0, fmt.Errorf("Error parsing ftp protocol: Status code too short: '%s'", response)
 	}
-	return status[0], nil
+	return strconv.Atoi(response[:3])
 }

starttls/starttls.go

@@ -34,6 +34,9 @@ import (
 	"github.com/mwitkow/go-http-dialer"
 )
 
+// Protocols are the names of supported protocols
+var Protocols []string = []string{"mysql", "postgres", "psql", "smtp", "ldap", "ftp"}
+
 type connectResult struct {
 	state *tls.ConnectionState
 	err   error