From 49e4cbe228f83812cd4b9e226e5b89bf6d778e97 Mon Sep 17 00:00:00 2001
From: David Newhall II <david@sleepers.pro>
Date: Mon, 8 Mar 2021 01:04:12 -0800
Subject: [PATCH 1/2] add ability to pass in pinned certificates.

---
 core/unifi/types.go | 25 +++++++++++++++----
 core/unifi/unifi.go | 59 ++++++++++++++++++++++++++++++++++++++-------
 2 files changed, 70 insertions(+), 14 deletions(-)

diff --git a/core/unifi/types.go b/core/unifi/types.go
index 1413402a..8676f85b 100644
--- a/core/unifi/types.go
+++ b/core/unifi/types.go
@@ -47,7 +47,7 @@ const (
 // path returns the correct api path based on the new variable.
 // new is based on the unifi-controller output. is it new or old output?
 func (u *Unifi) path(path string) string {
-	if u.New {
+	if u.new {
 		if path == APILoginPath {
 			return APILoginPathNew
 		}
@@ -80,16 +80,17 @@ type Devices struct {
 }
 
 // Config is the data passed into our library. This configures things and allows
-// us to connect to a controller and write log messages.
+// us to connect to a controller and write log messages. Optional SSLCert is used
+// for ssl cert pinning; provide the content of a PEM to validate the server's cert.
 type Config struct {
 	User      string
 	Pass      string
 	URL       string
-	VerifySSL bool
-	New       bool
+	SSLCert   [][]byte
 	ErrorLog  Logger
 	DebugLog  Logger
 	Timeout   time.Duration // how long to wait for replies, default: forever.
+	VerifySSL bool
 }
 
 // Unifi is what you get in return for providing a password! Unifi represents
@@ -100,7 +101,21 @@ type Unifi struct {
 	*http.Client
 	*Config
 	*server
-	csrf string
+	csrf         string
+	fingerprints fingerprints
+	new          bool
+}
+
+type fingerprints []string
+
+func (f fingerprints) Contains(s string) bool {
+	for i := range f {
+		if s == f[i] {
+			return true
+		}
+	}
+
+	return false
 }
 
 // server is the /status endpoint from the Unifi controller.
diff --git a/core/unifi/unifi.go b/core/unifi/unifi.go
index 4f05463b..93962065 100644
--- a/core/unifi/unifi.go
+++ b/core/unifi/unifi.go
@@ -8,8 +8,11 @@ package unifi
 import (
 	"bytes"
 	"context"
+	"crypto/sha256"
 	"crypto/tls"
+	"crypto/x509"
 	"encoding/json"
+	"encoding/pem"
 	"fmt"
 	"io"
 	"io/ioutil"
@@ -26,6 +29,7 @@ var (
 	ErrAuthenticationFailed = fmt.Errorf("authentication failed")
 	ErrInvalidStatusCode    = fmt.Errorf("invalid status code from server")
 	ErrNoParams             = fmt.Errorf("requedted PUT with no parameters")
+	ErrInvalidSignature     = fmt.Errorf("certificate signature does not match")
 )
 
 // NewUnifi creates a http.Client with authenticated cookies.
@@ -37,6 +41,29 @@ func NewUnifi(config *Config) (*Unifi, error) {
 		return nil, fmt.Errorf("creating cookiejar: %w", err)
 	}
 
+	u := newUnifi(config, jar)
+
+	for i, cert := range config.SSLCert {
+		p, _ := pem.Decode(cert)
+		u.fingerprints[i] = fmt.Sprintf("%x", sha256.Sum256(p.Bytes))
+	}
+
+	if err := u.checkNewStyleAPI(); err != nil {
+		return u, err
+	}
+
+	if err := u.Login(); err != nil {
+		return u, err
+	}
+
+	if err := u.GetServerData(); err != nil {
+		return u, fmt.Errorf("unable to get server version: %w", err)
+	}
+
+	return u, nil
+}
+
+func newUnifi(config *Config, jar http.CookieJar) *Unifi {
 	config.URL = strings.TrimRight(config.URL, "/")
 
 	if config.ErrorLog == nil {
@@ -53,24 +80,38 @@ func NewUnifi(config *Config) (*Unifi, error) {
 			Timeout: config.Timeout,
 			Jar:     jar,
 			Transport: &http.Transport{
-				TLSClientConfig: &tls.Config{InsecureSkipVerify: !config.VerifySSL}, // nolint: gosec
+				TLSClientConfig: &tls.Config{
+					InsecureSkipVerify: !config.VerifySSL, // nolint: gosec
+				},
 			},
 		},
 	}
 
-	if err := u.checkNewStyleAPI(); err != nil {
-		return u, err
+	if len(config.SSLCert) > 0 {
+		u.fingerprints = make(fingerprints, len(config.SSLCert))
+		u.Client.Transport = &http.Transport{
+			TLSClientConfig: &tls.Config{
+				InsecureSkipVerify:    true, // nolint: gosec
+				VerifyPeerCertificate: u.verifyPeerCertificate,
+			},
+		}
 	}
 
-	if err := u.Login(); err != nil {
-		return u, err
+	return u
+}
+
+func (u *Unifi) verifyPeerCertificate(certs [][]byte, chains [][]*x509.Certificate) error {
+	if len(u.fingerprints) == 0 {
+		return nil
 	}
 
-	if err := u.GetServerData(); err != nil {
-		return u, fmt.Errorf("unable to get server version: %w", err)
+	for _, cert := range certs {
+		if u.fingerprints.Contains(fmt.Sprintf("%x", sha256.Sum256(cert))) {
+			return nil
+		}
 	}
 
-	return u, nil
+	return ErrInvalidSignature
 }
 
 // Login is a helper method. It can be called to grab a new authentication cookie.
@@ -144,7 +185,7 @@ func (u *Unifi) checkNewStyleAPI() error {
 
 	if resp.StatusCode == http.StatusOK {
 		// The new version returns a "200" for a / request.
-		u.New = true
+		u.new = true
 		u.DebugLog("Using NEW UniFi controller API paths for %s", req.URL)
 	}
 

From f520681b825134f3f587398bc7de64a5a8b5b6c2 Mon Sep 17 00:00:00 2001
From: David Newhall II <david@sleepers.pro>
Date: Mon, 8 Mar 2021 01:06:48 -0800
Subject: [PATCH 2/2] add comment

---
 core/unifi/types.go | 1 +
 1 file changed, 1 insertion(+)

diff --git a/core/unifi/types.go b/core/unifi/types.go
index 8676f85b..5ae99f88 100644
--- a/core/unifi/types.go
+++ b/core/unifi/types.go
@@ -108,6 +108,7 @@ type Unifi struct {
 
 type fingerprints []string
 
+// Contains returns true if the fingerprint is in the list.
 func (f fingerprints) Contains(s string) bool {
 	for i := range f {
 		if s == f[i] {