finish support for exporting protonmail emails

common.py

@@ -7,22 +7,34 @@ from selenium.common.exceptions import TimeoutException
 from selenium.common.exceptions import ElementNotInteractableException
 
 import json
+import os
 from types import NoneType
 
 class MailMessage:
-    def __init__(self, subject: str, sender: str, maintext: str):
+    def __init__(self, index: str, sender: str, subject: str, maintext: str, acted_upon: bool = True):
         self.subject = subject
+        self.index = index
         self.sender = sender
         self.maintext = maintext
+        self.acted_upon = acted_upon
+
+        self.attrs = "subject index sender maintext acted_upon"
+        self.attrs = self.attrs.split(" ")
 
     def __eq__(self, other):
-        return other == self.subject or other == self.sender or other == self.maintext
+        for attr in self.attrs:
+            if other == getattr(self, attr):
+                return True
+
+        return False
 
     def serialize(self):
         return {
             "subject": self.subject,
             "sender": self.sender,
+            "index": self.index,
             "maintext": self.maintext,
+            "acted_upon": self.acted_upon,
             "date": "", # fill this in later
         }
     
@@ -36,15 +48,21 @@ class MailMessage:
         else:
             parsed = json_data
         
-        return cls(json_data["subject"], json_data["sender"], json_data["maintext"])
+        return cls(
+            index = json_data["index"],
+            sender = json_data["sender"],
+            subject = json_data["subject"],
+            maintext = json_data["maintext"],
+            acted_upon = json_data["acted_upon"],
+        )
 
 class MailProvider:
     def __init__(self, username: str, password: str, cachefile: NoneType | str = None):
         self.username: str = username
         self.password: str = password
-        self.sessionmessages = []
         self.webdriver = None # TODO: fill in type information for this instance
         self.cachefile = cachefile
+        self.seen_messages = []
 
     def _get_webdriver(self):
         self.webdriver = webdriver.Chrome()
@@ -64,7 +82,7 @@ class MailProvider:
             if not self.cachefile:
                 self.cachefile = self.__class__.__qualname__
 
-            with open(self.cachefile, "r") as f:
+            with open("data/{}".format(self.cachefile), "r") as f:
                 loaded = json.load(f)["payload"]
                 for entry in loaded:
                     self.seen_messages.append(MailMessage.from_json(entry))
@@ -78,9 +96,11 @@ class MailProvider:
         if not self.cachefile:
             self.cachefile = self.__name__
 
-        with open(self.cachefile, "w") as f:
+        if not os.path.exists("data/"):
+            os.mkdir("data/")
+
+        with open("data/{}".format(self.cachefile), "w") as f:
             data = {"payload": [item.serialize() for item in self.seen_messages]}
-            print(data)
             f.write(
                 json.dumps(
                     data

driver.py

@@ -1,4 +1,4 @@
 from proton import ProtonWebmail
 from secrets import proton, gmail
 
-print(ProtonWebmail(proton.username, proton.password).get())
+ProtonWebmail(proton.username, proton.password).get()

proton.py

@@ -5,11 +5,6 @@ import json
 from time import sleep
 
 class ProtonWebmail(MailProvider):
-    def transform_message_header(self, header):
-        return header.replace('\n', ' - ').replace('Unread email - ', '')
-
-    def get(self):
-        self.messages_failed = 0
     xpaths = {
         "username_box": "/html/body/div[1]/div[4]/div[1]/main/div[1]/div[2]/form/div[2]/div[1]/div/div/input",
         "password_box": "/html/body/div[1]/div[4]/div[1]/main/div[1]/div[2]/form/div[3]/div[1]/div/div[1]/input",
@@ -18,42 +13,52 @@ class ProtonWebmail(MailProvider):
         "messagebody": "/html/body/div[1]/div[3]/div/div[2]/div/div[2]/div/div/div/main/div/div/section/div/div[3]/div/div/article/div[2]",
         "backbutton": "//*[text()[contains(., 'Back')]]",
     }
+
+    def transform_message_header(self, header):
+        return header.replace('\n', ' - ').replace('Unread email - ', '')
+
+    def get(self):
+        self.messages_failed = 0
         self.get_seen_messages()
         self._get_webdriver()
         self.webdriver.get("https://account.proton.me/login")
-        self._wait_for_elem(xpaths["username_box"])
+        self._wait_for_elem(self.xpaths["username_box"])
         sleep(0.5)
-        self._type_in_elem(xpaths["username_box"], self.username)
+        self._type_in_elem(self.xpaths["username_box"], self.username)
-        self._click_elem(xpaths["password_box"])
+        self._click_elem(self.xpaths["password_box"])
-        self._type_in_elem(xpaths["password_box"], self.password)
+        self._type_in_elem(self.xpaths["password_box"], self.password)
-        self._click_elem(xpaths["sign_in"])
+        self._click_elem(self.xpaths["sign_in"])
-        self._wait_for_elem(xpaths["messages"])
+        self._wait_for_elem(self.xpaths["messages"])
-        count = 0
+        for i in self._get_elem_children(self._to_elem(self.xpaths["messages"])):
-        for i in self._get_elem_children(self._to_elem(xpaths["messages"])):
+            index = self.transform_message_header(i.text)
-            
+            if not self.is_seen(index):
-            if not self.is_seen(i.text):
-                count += 1
-                if count == 3:
-                    break
                 text = i.text
                 # we can interact with it, just selenium doesn't like it
                 try:
                     self._click_elem(i)
                 except ElementNotInteractableException:
                     pass
-                self._wait_for_elem(xpaths["messagebody"])
+                self._wait_for_elem(self.xpaths["messagebody"])
                 sleep(5)
-                self.webdriver.switch_to.frame(self._to_elem(xpaths["messagebody"]).find_elements(By.XPATH, "//iframe")[0])
+                self.webdriver.switch_to.frame(
+                    self._to_elem(self.xpaths["messagebody"]) \
+                        .find_elements(By.XPATH, "//iframe")[0]
+                )
+
+                splitted = index.split(' - ')
                 message = MailMessage(
-                    self.transform_message_header(text),
+                    index = self.transform_message_header(text),
-                    "",
+                    sender = splitted[0],
-                    self.webdriver.page_source,
+                    subject = splitted[3],
+                    maintext = self.webdriver.page_source,
                 )
                 self.add_to_seen(message)
                 self.webdriver.switch_to.default_content()
                 sleep(2)
-                self._click_elem(xpaths["backbutton"])
+                self._click_elem(self.xpaths["backbutton"])
                 sleep(2)
 
         self.write_seen_messages()
         self.webdriver.quit()
+
+        return self.seen_messages