From c6b479b01a91786666ded79a4cc126a3c2ddbfec Mon Sep 17 00:00:00 2001
From: emersion <contact@emersion.fr>
Date: Wed, 20 Sep 2017 10:09:31 +0200
Subject: [PATCH] protonmail: add some messages + public keys endpoints

---
 auth/auth.go           |  3 +-
 protonmail/keys.go     | 61 +++++++++++++++++++++++++++++++
 protonmail/messages.go | 82 ++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 145 insertions(+), 1 deletion(-)

diff --git a/auth/auth.go b/auth/auth.go
index c3cc5d5..5132051 100644
--- a/auth/auth.go
+++ b/auth/auth.go
@@ -8,10 +8,11 @@ import (
 	"io"
 	"os"
 
-	"github.com/emersion/hydroxide/protonmail"
 	"golang.org/x/crypto/bcrypt"
 	"golang.org/x/crypto/nacl/secretbox"
 	"golang.org/x/crypto/openpgp"
+
+	"github.com/emersion/hydroxide/protonmail"
 )
 
 const authFile = "auth.json"
diff --git a/protonmail/keys.go b/protonmail/keys.go
index 670ac8e..400c7d2 100644
--- a/protonmail/keys.go
+++ b/protonmail/keys.go
@@ -1,5 +1,14 @@
 package protonmail
 
+import (
+	"errors"
+	"net/http"
+	"net/url"
+	"strings"
+
+	"golang.org/x/crypto/openpgp"
+)
+
 type Key struct {
 	ID          string
 	Version     int
@@ -8,3 +17,55 @@ type Key struct {
 	Fingerprint string
 	Activation  interface{} // TODO
 }
+
+type RecipientType int
+
+const (
+	RecipientInternal RecipientType = 1
+	RecipientExternal               = 2
+)
+
+type PublicKeyResp struct {
+	RecipientType RecipientType
+	MIMEType      string
+	Keys          []*PublicKey
+}
+
+type PublicKey struct {
+	Send      int
+	PublicKey string
+}
+
+func (pub *PublicKey) Entity() (*openpgp.Entity, error) {
+	keyRing, err := openpgp.ReadArmoredKeyRing(strings.NewReader(pub.PublicKey))
+	if err != nil {
+		return nil, err
+	}
+	if len(keyRing) == 0 {
+		return nil, errors.New("public key is empty")
+	}
+	return keyRing[0], nil
+}
+
+// GetPublicKeys retrieves public keys for a user. The first key in
+// PublicKeyResp.Keys can be used for sending.
+func (c *Client) GetPublicKeys(email string) (*PublicKeyResp, error) {
+	v := url.Values{}
+	v.Set("Email", email)
+	// TODO: Fingerprint
+
+	req, err := c.newRequest(http.MethodGet, "/keys?"+v.Encode(), nil)
+	if err != nil {
+		return nil, err
+	}
+
+	var respData struct {
+		resp
+		*PublicKeyResp
+	}
+	if err := c.doJSON(req, &respData); err != nil {
+		return nil, err
+	}
+
+	return respData.PublicKeyResp, nil
+}
diff --git a/protonmail/messages.go b/protonmail/messages.go
index 8b5476d..9030d8a 100644
--- a/protonmail/messages.go
+++ b/protonmail/messages.go
@@ -1,7 +1,13 @@
 package protonmail
 
 import (
+	"bytes"
+	"io"
 	"net/http"
+	"strings"
+
+	"golang.org/x/crypto/openpgp"
+	"golang.org/x/crypto/openpgp/armor"
 )
 
 type MessageType int
@@ -15,6 +21,20 @@ const (
 
 type MessageEncryption int
 
+const (
+	MessageUnencrypted MessageEncryption = iota
+	MessageEncryptedInternal
+	MessageEncryptedExternal
+	MessageEncryptedOutside
+	_
+	_
+	_
+	MessageEncryptedInlinePGP
+	MessageEncryptedPGPMIME
+	_
+	_
+)
+
 type MessageAddress struct {
 	Address string
 	Name    string
@@ -48,6 +68,66 @@ type Message struct {
 	ExternalID     string
 }
 
+func (msg *Message) Read(keyring openpgp.KeyRing, prompt openpgp.PromptFunction) (*openpgp.MessageDetails, error) {
+	switch msg.IsEncrypted {
+	case MessageUnencrypted:
+		return &openpgp.MessageDetails{
+			IsEncrypted:    false,
+			IsSigned:       false,
+			UnverifiedBody: strings.NewReader(msg.Body),
+		}, nil
+	default:
+		block, err := armor.Decode(strings.NewReader(msg.Body))
+		if err != nil {
+			return nil, err
+		}
+
+		return openpgp.ReadMessage(block.Body, keyring, prompt, nil)
+	}
+}
+
+type messageWriter struct {
+	plaintext  io.WriteCloser
+	ciphertext io.WriteCloser
+	b          *bytes.Buffer
+	msg        *Message
+}
+
+func (w *messageWriter) Write(p []byte) (n int, err error) {
+	return w.plaintext.Write(p)
+}
+
+func (w *messageWriter) Close() error {
+	if err := w.plaintext.Close(); err != nil {
+		return err
+	}
+	if err := w.ciphertext.Close(); err != nil {
+		return err
+	}
+	w.msg.Body = w.b.String()
+	return nil
+}
+
+func (msg *Message) Encrypt(to []*openpgp.Entity, signed *openpgp.Entity) (plaintext io.WriteCloser, err error) {
+	var b bytes.Buffer
+	ciphertext, err := armor.Encode(&b, "PGP MESSAGE", nil)
+	if err != nil {
+		return nil, err
+	}
+
+	plaintext, err = openpgp.Encrypt(ciphertext, to, signed, nil, nil)
+	if err != nil {
+		return nil, err
+	}
+
+	return &messageWriter{
+		plaintext:  plaintext,
+		ciphertext: ciphertext,
+		b:          &b,
+		msg:        msg,
+	}, nil
+}
+
 func (c *Client) GetMessage(id string) (*Message, error) {
 	req, err := c.newRequest(http.MethodGet, "/messages/"+id, nil)
 	if err != nil {
@@ -65,6 +145,8 @@ func (c *Client) GetMessage(id string) (*Message, error) {
 	return respData.Message, nil
 }
 
+// CreateDraftMessage creates a new draft message. ToList, CCList, BCCList,
+// Subject, Body and AddressID are required in msg.
 func (c *Client) CreateDraftMessage(msg *Message, parentID string) (*Message, error) {
 	reqData := struct {
 		Message  *Message