#199 test fix for uploaded files

yb-support-tool/cmd/root.go

@@ -159,8 +159,14 @@ func Upload(email string, caseNum int, files []string) error {
 
 	}
 
-	fmt.Printf("File uploads complete\nFinalizing Package...\n")
+	fmt.Printf("File uploads complete\n")
 
+	fmt.Println("Uploading keycodes...")
+	if err := p.UploadKeycodes(); err != nil {
+		return fmt.Errorf("Unable to upload keycodes: %s", err)
+	}
+
+	fmt.Println("Finalizing Package...")
 	err = p.FinalizePackage()
 	if err != nil {
 		return fmt.Errorf("Unable to finalize package: %s", err)

yb-support-tool/sendsafelyuploader/encrypt.go

@@ -11,6 +11,7 @@ import (
 	"io"
 
 	"github.com/ProtonMail/go-crypto/openpgp"
+	"github.com/ProtonMail/go-crypto/openpgp/armor"
 	"github.com/ProtonMail/go-crypto/openpgp/packet"
 	"golang.org/x/crypto/pbkdf2"
 )
@@ -63,3 +64,32 @@ func (p *Package) Encrypt(unencryptedFilePart io.Reader) (*bytes.Buffer, error)
 	passphrase := p.Info.ServerSecret + p.Uploader.ClientSecret
 	return encrypt([]byte(passphrase), unencryptedFilePart)
 }
+
+func EncryptKeycode(keyCode string, publicKeyArmored string) (string, error) {
+	keyRing, err := openpgp.ReadArmoredKeyRing(bytes.NewBufferString(publicKeyArmored))
+	if err != nil {
+		return "", fmt.Errorf("failed to read public key: %w", err)
+	}
+
+	var encryptConfig packet.Config
+	encryptConfig.DefaultCipher = packet.CipherAES256
+	encryptConfig.DefaultHash = crypto.SHA256
+
+	buf := new(bytes.Buffer)
+	armorWriter, err := armor.Encode(buf, "PGP MESSAGE", nil)
+	if err != nil {
+		return "", fmt.Errorf("failed to create armor encoder: %w", err)
+	}
+
+	w, err := openpgp.Encrypt(armorWriter, keyRing, nil, nil, &encryptConfig)
+	if err != nil {
+		return "", fmt.Errorf("failed to create PGP encryptor: %w", err)
+	}
+	if _, err := io.WriteString(w, keyCode); err != nil {
+		return "", fmt.Errorf("failed to write keycode: %w", err)
+	}
+	w.Close()
+	armorWriter.Close()
+
+	return buf.String(), nil
+}

yb-support-tool/sendsafelyuploader/misc.go

@@ -14,11 +14,27 @@ type SSResponse struct {
 }
 
 type FinalizeInfo struct {
-	NeedsLink bool `json:"needsLink"`
+	NeedsLink  bool        `json:"needsLink"`
+	Recipients []Recipient `json:"recipients"`
 
 	SSResponse
 }
 
+type PublicKey struct {
+	ID  string `json:"id"`
+	Key string `json:"key"`
+}
+
+type PublicKeysResponse struct {
+	PublicKeys []PublicKey `json:"publicKeys"`
+	SSResponse
+}
+
+type Recipient struct {
+	RecipientID string `json:"recipientId"`
+	Email       string `json:"email"`
+}
+
 type HostedDropzoneInfo struct {
 	Success         string   `json:"success"`
 	Data            string   `json:"data"`

yb-support-tool/sendsafelyuploader/package.go

@@ -316,6 +316,81 @@ func (p *Package) MarkFileComplete(file *File) error {
 	return nil
 }
 
+func (p *Package) GetPublicKeys() (*PublicKeysResponse, error) {
+	endpoint := fmt.Sprintf("/drop-zone/v2.0/package/%s/public-keys/", p.Info.PackageCode)
+
+	body, err := p.Uploader.sendRequest(http.MethodGet, endpoint, nil)
+	if err != nil {
+		return nil, err
+	}
+
+	var resp PublicKeysResponse
+	if err := json.Unmarshal(body, &resp); err != nil {
+		return nil, err
+	}
+	if resp.Response != "SUCCESS" {
+		return nil, fmt.Errorf("failed to get public keys: %s", resp.Message)
+	}
+	return &resp, nil
+}
+
+func (p *Package) UploadKeycode(publicKeyID, encryptedKeycode string) error {
+	endpoint := fmt.Sprintf("/drop-zone/v2.0/package/%s/link/%s/", p.Info.PackageCode, publicKeyID)
+
+	reqBody := map[string]string{
+		"keycode": encryptedKeycode,
+	}
+
+	var reqBuf bytes.Buffer
+	enc := json.NewEncoder(&reqBuf)
+	enc.SetEscapeHTML(false)
+	if err := enc.Encode(reqBody); err != nil {
+		return err
+	}
+
+	body, err := p.Uploader.sendRequest(http.MethodPut, endpoint, reqBuf.Bytes(), withReqJSONHeader())
+	if err != nil {
+		return err
+	}
+
+	var resp SSResponse
+	if err := json.Unmarshal(body, &resp); err != nil {
+		return err
+	}
+	if resp.Response != "SUCCESS" {
+		return fmt.Errorf("%s", resp.Message)
+	}
+	return nil
+}
+
+func (p *Package) UploadKeycodes() error {
+	pubKeys, err := p.GetPublicKeys()
+	if err != nil {
+		return fmt.Errorf("failed to get public keys: %w", err)
+	}
+
+	var lastErr error
+	for _, pk := range pubKeys.PublicKeys {
+		encrypted, err := EncryptKeycode(p.Uploader.ClientSecret, pk.Key)
+		if err != nil {
+			lastErr = fmt.Errorf("failed to encrypt keycode for key '%s': %w", pk.ID, err)
+			fmt.Fprintf(os.Stderr, "Warning: %v\n", lastErr)
+			continue
+		}
+
+		if err := p.UploadKeycode(pk.ID, encrypted); err != nil {
+			lastErr = fmt.Errorf("failed to upload keycode for key '%s': %w", pk.ID, err)
+			fmt.Fprintf(os.Stderr, "Warning: %v\n", lastErr)
+			continue
+		}
+	}
+
+	if lastErr != nil {
+		return lastErr
+	}
+	return nil
+}
+
 // FinalizePackage marks the package complete and returns a link to the package,
 // or the error if package finalization did not succeed. Sets Package.URL on success
 func (p *Package) FinalizePackage() error {