From 6554458a4a15d1b37d5bdfe3625cb1bb7d4a01a1 Mon Sep 17 00:00:00 2001
From: Dustin Butler <dustin@fivetechnology.com>
Date: Sat, 22 Oct 2022 20:22:03 -0600
Subject: [PATCH] Implement alternative security

---
 conn.go            |   7 +++
 notbad_security.go | 144 +++++++++++++++++++++++++++++++++++++++++++++
 security.go        |   8 +++
 3 files changed, 159 insertions(+)
 create mode 100644 notbad_security.go

diff --git a/conn.go b/conn.go
index 2da23bd..67dc4d7 100644
--- a/conn.go
+++ b/conn.go
@@ -329,6 +329,13 @@ func (c *Conn) sendMulti(msg Msg) error {
 }
 
 func (c *Conn) send(isCommand bool, body []byte, flag byte) error {
+	// Encrypt body
+	body, err := c.sec.EncryptBody(body); 
+	if err != nil {
+		c.checkIO(err)
+		return err
+	}
+	
 	// Long flag
 	size := len(body)
 	isLong := size > 255
diff --git a/notbad_security.go b/notbad_security.go
new file mode 100644
index 0000000..b61099d
--- /dev/null
+++ b/notbad_security.go
@@ -0,0 +1,144 @@
+package zmq4
+
+import (
+	"fmt"
+	"io"
+	"log"
+
+	ecies "github.com/ecies/go/v2"
+)
+
+
+type BadSecurity struct {}
+
+// Type returns the security mechanism type.
+func (BadSecurity) Type() SecurityType {
+	return "NOTBAD"
+}
+
+// Handshake implements the ZMTP security handshake according to
+// this security mechanism.
+// see:
+//
+//	https://rfc.zeromq.org/spec:23/ZMTP/
+//	https://rfc.zeromq.org/spec:24/ZMTP-PLAIN/
+//	https://rfc.zeromq.org/spec:25/ZMTP-CURVE/
+
+func handshakeServer(conn *Conn) error {
+
+	raw, err := conn.Meta.MarshalZMTP()
+	if err != nil {
+		return fmt.Errorf("zmq4: could not marshal metadata: %w", err)
+	}
+
+	cmd, err := conn.RecvCmd()
+	log.Printf("Got %s", cmd.Name)
+	if err != nil {
+		return fmt.Errorf("zmq4: could not recv metadata from peer: %w", err)
+	}
+	if cmd.Name != CmdHello {
+		return ErrBadCmd
+	}
+	
+	log.Printf("Send %s", CmdWelcome)
+	err = conn.SendCmd(CmdWelcome, []byte{})
+	if err != nil {
+		return fmt.Errorf("zmq4: could not send metadata to peer: %w", err)
+	}
+
+	log.Printf("Send %s", CmdReady)
+	err = conn.SendCmd(CmdReady, raw)
+	if err != nil {
+		return fmt.Errorf("zmq4: could not send metadata to peer: %w", err)
+	}
+
+	cmd, err = conn.RecvCmd()
+	log.Printf("Got %s", cmd.Name)
+	if err != nil {
+		return fmt.Errorf("zmq4: could not recv metadata from peer: %w", err)
+	}
+
+	if cmd.Name != CmdReady {
+		return ErrBadCmd
+	}
+
+	err = conn.Peer.Meta.UnmarshalZMTP(cmd.Body)
+	if err != nil {
+		return fmt.Errorf("zmq4: could not unmarshal peer metadata: %w", err)
+	}
+	log.Print("Server handshake complete")
+	return nil
+}
+
+func (BadSecurity) Handshake(conn *Conn, server bool) error {
+	if server {
+		return handshakeServer(conn)
+	}
+	
+	raw, err := conn.Meta.MarshalZMTP()
+	if err != nil {
+		return fmt.Errorf("zmq4: could not marshal metadata: %w", err)
+	}
+
+	log.Printf("Send %s", CmdHello)
+	err = conn.SendCmd(CmdHello, []byte{})
+	if err != nil {
+		return fmt.Errorf("zmq4: could not send metadata to peer: %w", err)
+	}
+	
+	cmd, err := conn.RecvCmd()
+	log.Printf("Got %s", cmd.Name)
+	if err != nil {
+		return fmt.Errorf("zmq4: could not recv metadata from peer: %w", err)
+	}
+	if cmd.Name != CmdWelcome {
+		return ErrBadCmd
+	}
+
+	log.Printf("Send %s", CmdReady)
+	err = conn.SendCmd(CmdReady, raw)
+	if err != nil {
+		return fmt.Errorf("zmq4: could not send metadata to peer: %w", err)
+	}
+
+	cmd, err = conn.RecvCmd()
+	log.Printf("Got %s", cmd.Name)
+	if err != nil {
+		return fmt.Errorf("zmq4: could not recv metadata from peer: %w", err)
+	}
+
+	if cmd.Name != CmdReady {
+		return ErrBadCmd
+	}
+
+	err = conn.Peer.Meta.UnmarshalZMTP(cmd.Body)
+	if err != nil {
+		return fmt.Errorf("zmq4: could not unmarshal peer metadata: %w", err)
+	}
+
+	log.Print("Client handshake complete")
+	return nil
+}
+
+// Encrypt body
+func (BadSecurity) EncryptBody(data []byte) ([]byte, error) {
+	key, _ := ecies.NewPrivateKeyFromHex("14533aad6a633f2a814e23700637540573aba715f94b519b1c01219d645153d7")
+	encrypted, _ := ecies.Encrypt(key.PublicKey, data)
+	return encrypted, nil
+}
+
+// Encrypt writes the encrypted form of data to w.
+func (BadSecurity) Encrypt(w io.Writer, data []byte) (int, error) {
+	return w.Write(data)
+}
+
+// Decrypt writes the decrypted form of data to w.
+func (BadSecurity) Decrypt(w io.Writer, data []byte) (int, error) {
+	key, _ := ecies.NewPrivateKeyFromHex("14533aad6a633f2a814e23700637540573aba715f94b519b1c01219d645153d7")
+	decrypted, err := ecies.Decrypt(key, data)
+	if err != nil {
+		log.Printf("Decrypt Error: %v %s", data, data)
+		return w.Write(data)
+	}
+	return w.Write(decrypted)
+}
\ No newline at end of file
diff --git a/security.go b/security.go
index 959df89..292c174 100644
--- a/security.go
+++ b/security.go
@@ -25,6 +25,9 @@ type Security interface {
 	// Encrypt writes the encrypted form of data to w.
 	Encrypt(w io.Writer, data []byte) (int, error)
 
+	// Encrypt body before send calculates headers
+	EncryptBody(data []byte) ([]byte, error)
+
 	// Decrypt writes the decrypted form of data to w.
 	Decrypt(w io.Writer, data []byte) (int, error)
 }
@@ -95,6 +98,11 @@ func (nullSecurity) Encrypt(w io.Writer, data []byte) (int, error) {
 	return w.Write(data)
 }
 
+// Encrypt body before send calculates headers.
+func (nullSecurity) EncryptBody(data []byte) ([]byte, error) {
+	return data, nil
+}
+
 // Decrypt writes the decrypted form of data to w.
 func (nullSecurity) Decrypt(w io.Writer, data []byte) (int, error) {
 	return w.Write(data)