From f3b4fc742c3213a82d60d3d87d6d7a709a6104e1 Mon Sep 17 00:00:00 2001
From: lqqyt2423 <974923609@qq.com>
Date: Thu, 3 Dec 2020 22:31:23 +0800
Subject: [PATCH] cert create

---
 proxy/cert.go      | 89 ++++++++++++++++++++++++++++++++++++++++++----
 proxy/cert_test.go | 26 ++++++++++----
 2 files changed, 102 insertions(+), 13 deletions(-)

diff --git a/proxy/cert.go b/proxy/cert.go
index 3edb375..a15b338 100644
--- a/proxy/cert.go
+++ b/proxy/cert.go
@@ -1,15 +1,20 @@
 package proxy
 
 import (
+	"crypto/rand"
 	"crypto/rsa"
 	"crypto/x509"
+	"crypto/x509/pkix"
 	"encoding/pem"
 	"errors"
 	"fmt"
 	"io"
 	"io/ioutil"
+	"log"
+	"math/big"
 	"os"
 	"path/filepath"
+	"time"
 )
 
 var caErrNotFound = errors.New("ca not found")
@@ -32,8 +37,11 @@ func NewCA(path string) (*CA, error) {
 		if err != caErrNotFound {
 			return nil, err
 		}
+	} else {
+		return ca, nil
 	}
 
+	log.Println("begin create ca")
 	if err := ca.create(); err != nil {
 		return nil, err
 	}
@@ -47,7 +55,7 @@ func getStorePath(path string) (string, error) {
 		if err != nil {
 			return "", err
 		}
-		path = filepath.Join(homeDir, ".mitmproxy")
+		path = filepath.Join(homeDir, ".go_mitmproxy")
 	}
 
 	if !filepath.IsAbs(path) {
@@ -77,8 +85,18 @@ func getStorePath(path string) (string, error) {
 	return path, nil
 }
 
+// The certificate and the private key in PEM format.
+func (ca *CA) caFile() string {
+	return filepath.Join(ca.StorePath, "mitmproxy-ca.pem")
+}
+
+// The certificate in PEM format.
+func (ca *CA) caCertFile() string {
+	return filepath.Join(ca.StorePath, "mitmproxy-ca-cert.pem")
+}
+
 func (ca *CA) load() error {
-	caFile := filepath.Join(ca.StorePath, "mitmproxy-ca.pem")
+	caFile := ca.caFile()
 	stat, err := os.Stat(caFile)
 	if err != nil {
 		if os.IsNotExist(err) {
@@ -125,7 +143,50 @@ func (ca *CA) load() error {
 }
 
 func (ca *CA) create() error {
-	return nil
+	key, err := rsa.GenerateKey(rand.Reader, 2048)
+	if err != nil {
+		return err
+	}
+	ca.PrivateKey = *key
+
+	template := &x509.Certificate{
+		SerialNumber: big.NewInt(time.Now().UnixNano() / 100000),
+		Subject: pkix.Name{
+			CommonName:   "mitmproxy",
+			Organization: []string{"mitmproxy"},
+		},
+		NotBefore:             time.Now().Add(-time.Hour * 48),
+		NotAfter:              time.Now().Add(time.Hour * 24 * 365 * 3),
+		BasicConstraintsValid: true,
+		IsCA:                  true,
+		SignatureAlgorithm:    x509.SHA256WithRSA,
+		KeyUsage:              x509.KeyUsageCertSign | x509.KeyUsageCRLSign,
+		ExtKeyUsage: []x509.ExtKeyUsage{
+			x509.ExtKeyUsageServerAuth,
+			x509.ExtKeyUsageClientAuth,
+			x509.ExtKeyUsageEmailProtection,
+			x509.ExtKeyUsageTimeStamping,
+			x509.ExtKeyUsageCodeSigning,
+			x509.ExtKeyUsageMicrosoftCommercialCodeSigning,
+			x509.ExtKeyUsageMicrosoftServerGatedCrypto,
+			x509.ExtKeyUsageNetscapeServerGatedCrypto,
+		},
+	}
+
+	certBytes, err := x509.CreateCertificate(rand.Reader, template, template, &key.PublicKey, key)
+	if err != nil {
+		return err
+	}
+	cert, err := x509.ParseCertificate(certBytes)
+	if err != nil {
+		return err
+	}
+	ca.RootCert = *cert
+
+	if err := ca.save(); err != nil {
+		return err
+	}
+	return ca.saveCert()
 }
 
 func (ca *CA) saveTo(out io.Writer) error {
@@ -138,15 +199,29 @@ func (ca *CA) saveTo(out io.Writer) error {
 		return err
 	}
 
-	err = pem.Encode(out, &pem.Block{Type: "CERTIFICATE", Bytes: ca.RootCert.Raw})
+	return pem.Encode(out, &pem.Block{Type: "CERTIFICATE", Bytes: ca.RootCert.Raw})
+}
+
+func (ca *CA) saveCertTo(out io.Writer) error {
+	return pem.Encode(out, &pem.Block{Type: "CERTIFICATE", Bytes: ca.RootCert.Raw})
+}
+
+func (ca *CA) save() error {
+	file, err := os.Create(ca.caFile())
 	if err != nil {
 		return err
 	}
+	defer file.Close()
 
-	return nil
+	return ca.saveTo(file)
 }
 
-func (ca *CA) save() error {
+func (ca *CA) saveCert() error {
+	file, err := os.Create(ca.caCertFile())
+	if err != nil {
+		return err
+	}
+	defer file.Close()
 
-	return nil
+	return ca.saveCertTo(file)
 }
diff --git a/proxy/cert_test.go b/proxy/cert_test.go
index 69a3d55..d543855 100644
--- a/proxy/cert_test.go
+++ b/proxy/cert_test.go
@@ -1,28 +1,42 @@
 package proxy
 
 import (
-	"os"
+	"bytes"
+	"io/ioutil"
+	"reflect"
 	"testing"
 )
 
 func TestGetStorePath(t *testing.T) {
 	path, err := getStorePath("")
 	if err != nil {
-		t.Error(err)
+		t.Fatal(err)
 	}
 	if path == "" {
-		t.Errorf("should have path")
+		t.Fatal("should have path")
 	}
 }
 
 func TestNewCA(t *testing.T) {
 	ca, err := NewCA("")
 	if err != nil {
-		t.Error(err)
+		t.Fatal(err)
 	}
 
-	err = ca.saveTo(os.Stdout)
+	data := make([]byte, 0)
+	buf := bytes.NewBuffer(data)
+
+	err = ca.saveTo(buf)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	fileContent, err := ioutil.ReadFile(ca.caFile())
 	if err != nil {
-		t.Error(err)
+		t.Fatal(err)
+	}
+
+	if !reflect.DeepEqual(fileContent, buf.Bytes()) {
+		t.Fatal("pem content should equal")
 	}
 }