From ce009228ac3261db772c4ae17ed9e94cb9dfbddf Mon Sep 17 00:00:00 2001
From: Cedric Staub <cs@squareup.com>
Date: Fri, 3 Jun 2016 12:14:46 -0700
Subject: [PATCH] Guess more formats, clean up guessing code

---
 main.go | 42 ++++++++++++++++++++++++++++++------------
 1 file changed, 30 insertions(+), 12 deletions(-)

diff --git a/main.go b/main.go
index dc40325..583f960 100644
--- a/main.go
+++ b/main.go
@@ -23,6 +23,7 @@ import (
 	"fmt"
 	"io/ioutil"
 	"os"
+	"path/filepath"
 	"strings"
 
 	"golang.org/x/crypto/pkcs12"
@@ -34,17 +35,32 @@ var (
 
 	dump     = app.Command("dump", "Display information about a certificate.")
 	dumpFile = dump.Arg("file", "Certificate file to dump.").Required().String()
-	dumpType = dump.Flag("format", "Format of given input. If unspecified, certigo guesses based on file extension").Default("guess").Short('f').String()
+	dumpType = dump.Flag("format", "Format of given input. If unspecified, certigo guesses based on file extension").Short('f').String()
 )
 
+var fileExtToFormat = map[string]string{
+	".pem":   "PEM",
+	".crt":   "PEM",
+	".p12":   "PKCS12",
+	".pfx":   "PKCS12",
+	".jceks": "JCEKS",
+}
+
 func main() {
 	switch kingpin.MustParse(app.Parse(os.Args[1:])) {
-	case dump.FullCommand(): //Dump certificate
-		certs, err := getCerts(*dumpFile, *dumpType)
+	case dump.FullCommand(): // Dump certificate
+		format, ok := formatForFile(*dumpFile, *dumpType)
+		if !ok {
+			fmt.Fprint(os.Stderr, "unable to guess file type\n")
+			os.Exit(1)
+		}
+
+		certs, err := getCerts(*dumpFile, format)
 		if err != nil {
 			fmt.Fprintf(os.Stderr, "%s\n", err)
 			os.Exit(1)
 		}
+
 		for i, cert := range certs {
 			fmt.Println("CERTIFICATE", i+1)
 			displayCert(cert)
@@ -53,6 +69,16 @@ func main() {
 	}
 }
 
+// formatForFile returns the file format (either from flags or
+// based on file extension).
+func formatForFile(filename, format string) (string, bool) {
+	if format == "" {
+		guess, ok := fileExtToFormat[strings.ToLower(filepath.Ext(filename))]
+		return guess, ok
+	}
+	return format, true
+}
+
 // getCerts takes in a filename and format type and returns an
 // array of all the certificates found in that file. If no format
 // is specified for the file, getCerts guesses what format was used
@@ -89,16 +115,8 @@ func getCerts(file, format string) ([]*x509.Certificate, error) {
 				certs = append(certs, cert)
 			}
 		}
-	case "guess":
-		if strings.HasSuffix(file, "pem") {
-			return getCerts(file, "PEM")
-		} else if strings.HasSuffix(file, "p12") || strings.HasSuffix(file, "pks") {
-			return getCerts(file, "PKCS12")
-		} else {
-			return getCerts(file, ", couldn't guess format")
-		}
 	default:
-		return nil, fmt.Errorf("unknown type %s", format)
+		return nil, fmt.Errorf("unknown file type: %s", format)
 	}
 	return certs, nil
 }