Add support for Delta

Also add an undocumented mailbox driver for a directory of files
for ease of testing (this could probably become a maildir driver
with a bit more work).

Remove unecessary decode calls from the message traversal.

8 files changed, 152 insertions, 14 deletions

diff --git a/README.rst b/README.rst
index 8715aa6..dd22c95 100644
--- a/README.rst
+++ b/README.rst
@@ -14,9 +14,10 @@ even if an email is seen multiple times.
 
 It currently understands emails from the following senders:
 
-* United Airlines
-* Marriott hotels (when booked directly)
+* Delta Air Lines
 * Eventbrite
+* Marriott hotels (when booked directly)
+* United Airlines
 
 It is simple to add support for more types of emails, and additions
 are welcome.
diff --git a/email_assistant/assistant.py b/email_assistant/assistant.py
index a784077..d39b016 100644
--- a/email_assistant/assistant.py
+++ b/email_assistant/assistant.py
@@ -37,7 +37,7 @@ from email_assistant import plugins
 # Number of days to look backwards when scanning a mailbox for the first time:
 IMAP_BACKFILL = 180
 
-class Mailbox:
+class IMAPMailbox:
     def __init__(self, name, host, username, password, folders):
         self.log = logging.getLogger('assistant.mailbox')
         self.name = name
@@ -85,6 +85,21 @@ class Mailbox:
         with open(self.state_file, 'w') as f:
             json.dump(self.uidinfo, f)
 
+class DirMailbox:
+    def __init__(self, name, directory):
+        self.log = logging.getLogger('assistant.mailbox')
+        self.name = name
+        self.directory = directory
+
+    def get_messages(self):
+        for fn in os.listdir(self.directory):
+            with open(os.path.join(self.directory, fn), 'rb') as f:
+                msg = f.read()
+            yield msg
+
+    def save(self):
+        pass
+
 class Calendar:
     def __init__(self, url, username, password, calendar):
         self.log = logging.getLogger('assistant.calendar')
@@ -116,7 +131,7 @@ class Calendar:
 
 class Assistant:
     def __init__(self):
-        self.log = logging.getLogger('main')
+        self.log = logging.getLogger('assistant.main')
         self.geolocator = None
         self.tzfinder = None
         self.plugins = []
@@ -151,12 +166,17 @@ class Assistant:
         for section in config.sections():
             if section.startswith('mailbox '):
                 name = section.split()[1]
-                mailboxes[name] = Mailbox(
-                    name,
-                    config[section]['host'],
-                    config[section]['username'],
-                    config[section]['password'],
-                    config[section]['folders'].split(','))
+                if config[section]['type'].lower() == 'imap':
+                    mailboxes[name] = IMAPMailbox(
+                        name,
+                        config[section]['host'],
+                        config[section]['username'],
+                        config[section]['password'],
+                        config[section]['folders'].split(','))
+                elif config[section]['type'].lower() == 'dir':
+                    mailboxes[name] = DirMailbox(
+                        name,
+                        config[section]['path'])
             elif section.startswith('calendar '):
                 name = section.split()[1]
                 calendars[name] = Calendar(
diff --git a/email_assistant/plugins/__init__.py b/email_assistant/plugins/__init__.py
index 10b28f0..295852e 100644
--- a/email_assistant/plugins/__init__.py
+++ b/email_assistant/plugins/__init__.py
@@ -21,9 +21,11 @@
 from . import eventbrite
 from . import marriott
 from . import united
+from . import delta
 
 plugins = [
     eventbrite.Plugin,
     marriott.Plugin,
     united.Plugin,
+    delta.Plugin,
 ]
diff --git a/email_assistant/plugins/delta.py b/email_assistant/plugins/delta.py
new file mode 100644
index 0000000..50da506
--- /dev/null
+++ b/email_assistant/plugins/delta.py
@@ -0,0 +1,115 @@
+#!/usr/bin/env python3
+
+# Copyright (C) 2019 James E. Blair <corvus@gnu.org>
+#
+# This file is part of Email-assistant.
+#
+# Email-assistant is free software: you can redistribute it and/or
+# modify it under the terms of the GNU Affero General Public License
+# as published by the Free Software Foundation, either version 3 of
+# the License, or (at your option) any later version.
+#
+# Email-assistant is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
+# General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with Email-assistant.  If not, see
+# <https://www.gnu.org/licenses/>.
+
+import re
+import logging
+import hashlib
+
+from bs4 import BeautifulSoup
+import dateutil.parser
+import dateutil.tz
+import inscriptis
+import vobject
+
+from email_assistant import iata
+from email_assistant import plugin
+
+def parse_dep_arr(dep_date, flight_date, flight_time, code):
+    flight_year = dep_date.year
+    tz = iata.tzmap[code]
+    flight_time = dateutil.parser.parse(flight_date + ' ' + flight_time)
+    if flight_time.year < flight_year:
+        flight_time.replace(year=flight_year)
+    flight_time = flight_time.replace(tzinfo=dateutil.tz.gettz(tz))
+    return flight_time
+
+class Plugin(plugin.Plugin):
+    name = 'delta'
+
+    def match(self, msg):
+        if ('DeltaAirLines@e.delta.com' in msg['From'] and
+            'Your Flight Receipt' in msg['Subject']):
+            return True
+
+    def get_events(self, msg):
+        events = []
+        for part in msg.walk():
+            if part.get_content_type() == 'text/html':
+                soup = BeautifulSoup(part.get_payload(decode=True), 'html.parser')
+                flights = []
+                flight = None
+                dep_date = None
+                start = soup.find(string=re.compile('FLIGHT INFO STARTS'))
+                for row in start.parent.find_all('tr', recursive=False):
+                    row = [x.strip() for x in row.strings if x.strip()]
+                    if len(row) == 3 and row[1:] == ['DEPART', 'ARRIVE']:
+                        dep_date = row[0]
+                        continue
+                    if dep_date is None: continue
+                    f = {
+                        'dep_date': dep_date,
+                        #'num': row[0],
+                        #'cabin': row[1],
+                        'dep_city': row[2],
+                        'dep_time': row[3],
+                        'arr_city': row[4],
+                        'arr_time': row[5]}
+                    if len(row) > 6:
+                        f['arr_date'] = row[6].replace('*', '').strip()
+                    else:
+                        f['arr_date'] = dep_date
+                    flights.append(f)
+                start = soup.find(string=re.compile('Checked Bag Allowance'))
+                while start.name != 'table': start = start.parent
+                start = start.parent
+                while start.name != 'table': start = start.parent
+                index = 0
+                for row in start.find_all('tr', recursive=False):
+                    row = [x.strip() for x in row.strings if x.strip()]
+                    if len(row) != 3: continue
+                    try:
+                        dep_date = dateutil.parser.parse(row[0])
+                    except Exception:
+                        continue
+                    f = flights[index]
+                    f['dep_code'] = row[1].split()[-1].strip()
+                    f['arr_code'] = row[2].split()[-1].strip()
+                    f['dep_dt'] = parse_dep_arr(
+                        dep_date, f['dep_date'], f['dep_time'], f['dep_code'])
+                    f['arr_dt'] = parse_dep_arr(
+                        dep_date, f['arr_date'], f['arr_time'], f['arr_code'])
+                    index += 1
+                for f in flights:
+                    self.log.debug('%s %s %s %s',
+                                   f['dep_dt'], f['arr_dt'], f['dep_code'], f['arr_code'])
+                    cal = vobject.iCalendar()
+                    event = cal.add('vevent')
+                    event.add('dtstart').value = f['dep_dt']
+                    event.add('dtend').value = f['arr_dt']
+                    summary = "Flight from %s to %s" % (f['dep_code'], f['arr_code'])
+                    event.add('summary').value = summary
+                    text = inscriptis.get_text(str(soup))
+                    text = re.sub(r'([^ ]+)\s*\n', '\\1\n', text)
+                    event.add('description').value = text
+                    event.add('location').value = "%s airport" % (f['dep_code'],)
+                    uid = hashlib.sha1((str(f['dep_time']) + summary).encode('utf8')).hexdigest()
+                    event.add('uid').value = uid
+                    events.append(cal)
+        return events
diff --git a/email_assistant/plugins/eventbrite.py b/email_assistant/plugins/eventbrite.py
index 9ef073a..afd869f 100644
--- a/email_assistant/plugins/eventbrite.py
+++ b/email_assistant/plugins/eventbrite.py
@@ -44,7 +44,7 @@ class Plugin(plugin.Plugin):
         events = []
         for part in msg.walk():
             if part.get_content_type() == 'text/html':
-                soup = BeautifulSoup(part.get_payload(decode=True).decode('utf8'), 'html.parser')
+                soup = BeautifulSoup(part.get_payload(decode=True), 'html.parser')
                 data = json.loads(soup.find('script', type="application/ld+json").string)
 
                 address = data['reservationFor']['location']['address']
diff --git a/email_assistant/plugins/general.py b/email_assistant/plugins/general.py
index c0517b8..e21274d 100644
--- a/email_assistant/plugins/general.py
+++ b/email_assistant/plugins/general.py
@@ -24,7 +24,7 @@ def get_events(msg):
     log = logging.getLogger('assistant.marriott')
     for part in msg.walk():
         if part.get_content_type() == 'text/html':
-            soup = BeautifulSoup(part.get_payload(decode=True).decode('utf8'), 'html.parser')
+            soup = BeautifulSoup(part.get_payload(decode=True), 'html.parser')
 
             start = end = location = None
             for element in soup.descendants:
diff --git a/email_assistant/plugins/marriott.py b/email_assistant/plugins/marriott.py
index 52597cd..1f2abab 100644
--- a/email_assistant/plugins/marriott.py
+++ b/email_assistant/plugins/marriott.py
@@ -43,7 +43,7 @@ class Plugin(plugin.Plugin):
         events = []
         for part in msg.walk():
             if part.get_content_type() == 'text/html':
-                soup = BeautifulSoup(part.get_payload(decode=True).decode('utf8'), 'html.parser')
+                soup = BeautifulSoup(part.get_payload(decode=True), 'html.parser')
                 summary = soup.find_all('table')[7].a.string.strip()
                 location = soup.find_all('table')[9].a.string.strip()
                 start = (soup.find('th', string=re.compile('Check-In:')).
diff --git a/email_assistant/plugins/united.py b/email_assistant/plugins/united.py
index 0b33cd9..38de3e2 100644
--- a/email_assistant/plugins/united.py
+++ b/email_assistant/plugins/united.py
@@ -60,7 +60,7 @@ class Plugin(plugin.Plugin):
         events = []
         for part in msg.walk():
             if part.get_content_type() == 'text/html':
-                soup = BeautifulSoup(part.get_payload(decode=True).decode('utf8'), 'html.parser')
+                soup = BeautifulSoup(part.get_payload(decode=True), 'html.parser')
                 # confirmation_number = soup.find(class_="eTicketConfirmation").string
 
                 index = 0