From ea7f3ff2799aa50bbc88d4900d3dfe1895d26f65 Mon Sep 17 00:00:00 2001
From: Ross Smith II <ross@smithii.com>
Date: Thu, 21 Apr 2022 08:05:18 -0700
Subject: [PATCH] Add attachment headers

---
 simplegmail/attachment.py | 31 +++++++++++++++++--------------
 simplegmail/gmail.py      | 18 ++++++++++++++++--
 2 files changed, 33 insertions(+), 16 deletions(-)

diff --git a/simplegmail/attachment.py b/simplegmail/attachment.py
index 079f7c2..b31ae7c 100644
--- a/simplegmail/attachment.py
+++ b/simplegmail/attachment.py
@@ -11,7 +11,7 @@
 
 class Attachment(object):
     """
-    The Attachment class for attachments to emails in your Gmail mailbox. This 
+    The Attachment class for attachments to emails in your Gmail mailbox. This
     class should not be manually instantiated.
 
     Args:
@@ -22,6 +22,7 @@ class should not be manually instantiated.
         filename: The filename associated with the attachment.
         filetype: The mime type of the file.
         data: The raw data of the file. Default None.
+        headers: A dict of header name/value pairs of the attachment. Default {}.
 
     Attributes:
         _service (googleapiclient.discovery.Resource): The Gmail service object.
@@ -31,9 +32,10 @@ class should not be manually instantiated.
         filename (str): The filename associated with the attachment.
         filetype (str): The mime type of the file.
         data (bytes): The raw data of the file.
+        headers (dict): A dict of header name/value pairs of the attachment.
 
     """
-    
+
     def __init__(
         self,
         service: 'googleapiclient.discovery.Resource',
@@ -42,7 +44,8 @@ def __init__(
         att_id: str,
         filename: str,
         filetype: str,
-        data: Optional[bytes] = None
+        data: Optional[bytes] = None,
+        headers: Optional[dict] = {}
     ) -> None:
         self._service = service
         self.user_id = user_id
@@ -51,17 +54,18 @@ def __init__(
         self.filename = filename
         self.filetype = filetype
         self.data = data
+        self.headers = headers
 
     def download(self) -> None:
         """
         Downloads the data for an attachment if it does not exist.
-        
+
         Raises:
-            googleapiclient.errors.HttpError: There was an error executing the 
+            googleapiclient.errors.HttpError: There was an error executing the
                 HTTP request.
-        
+
         """
-        
+
         if self.data is not None:
             return
 
@@ -79,18 +83,18 @@ def save(
     ) -> None:
         """
         Saves the attachment. Downloads file data if not downloaded.
-        
+
         Args:
-            filepath: where to save the attachment. Default None, which uses 
+            filepath: where to save the attachment. Default None, which uses
                 the filename stored.
             overwrite: whether to overwrite existing files. Default False.
-        
+
         Raises:
-            FileExistsError: if the call would overwrite an existing file and 
+            FileExistsError: if the call would overwrite an existing file and
                 overwrite is not set to True.
-        
+
         """
-        
+
         if filepath is None:
             filepath = self.filename
 
@@ -105,4 +109,3 @@ def save(
 
         with open(filepath, 'wb') as f:
             f.write(self.data)
-
diff --git a/simplegmail/gmail.py b/simplegmail/gmail.py
index 2c9aba7..5b96aef 100644
--- a/simplegmail/gmail.py
+++ b/simplegmail/gmail.py
@@ -737,7 +737,7 @@ def _build_message_from_ref(
                 elif part['part_type'] == 'attachment':
                     attm = Attachment(self.service, user_id, msg_id,
                                       part['attachment_id'], part['filename'],
-                                      part['filetype'], part['data'])
+                                      part['filetype'], part['data'], part['headers'])
                     attms.append(attm)
 
             return Message(self.service, self.creds, user_id, msg_id,
@@ -782,12 +782,26 @@ def _evaluate_message_payload(
             if not filename:
                 filename = 'unknown'
 
+            headers = {}
+            for hdr in payload['headers']:
+                name = hdr['name']
+                value = hdr['value']
+                if name in headers:
+                    # Allow for duplicate header names, per page 20 of RFC 5322:
+                    # https://www.rfc-editor.org/rfc/rfc5322#section-3.6
+                    if not isinstance(headers[name], list):
+                        headers[name] = [headers[name]]
+                    headers[name].append(value)
+                    continue
+                headers[name] = value
+
             obj = {
                 'part_type': 'attachment',
                 'filetype': payload['mimeType'],
                 'filename': filename,
                 'attachment_id': att_id,
-                'data': None
+                'data': None,
+                'headers': headers
             }
 
             if attachments == 'reference':