Replace login system with flask-login.

2 files changed, 77 insertions, 49 deletions

diff --git a/main.py b/main.py
index 8315d3b..8bc64eb 100644
--- a/main.py
+++ b/main.py
@@ -4,6 +4,13 @@ from urllib.parse import urlencode
 import password
 from password import Passwords
 import os
+from flask_login import (
+    LoginManager,
+    login_required,
+    login_user,
+    current_user,
+    logout_user,
+)
 from typing import (
     Optional,
     cast,
@@ -22,6 +29,8 @@ from flask import (
     get_flashed_messages
 )
 
+login_manager = LoginManager()
+
 
 def mailto(addr: str) -> HTML:
     return ('a', {'href': f'mailto:{addr}'}, addr)
@@ -126,8 +135,7 @@ def flashed_messages() -> HTML:
 
 
 def page_base(title: Optional[str] = None,
-              body: HTML = [],
-              username: Optional[str] = None) -> HTML:
+              body: HTML = []) -> HTML:
     return ('html',
             ('head',
              ('meta', {'charset': 'utf-8'}),
@@ -136,7 +144,7 @@ def page_base(title: Optional[str] = None,
              ),
             ('body',
              ('nav',
-              user_info(username) if username else login_prompt()
+              user_info(current_user.get_id()) if current_user.is_authenticated else login_prompt()
               ),
              flashed_messages(),
              body))
@@ -179,8 +187,7 @@ def response_for(id: str, username: Optional[str] = None) -> str:
                  ('ul', attachement_tree(mail)),
                  ]
     html_str = render_document(page_base(title=title,
-                                         body=main_body,
-                                         username=username))
+                                         body=main_body))
 
     return html_str
 
@@ -194,39 +201,46 @@ def search_field(q: str) -> HTML:
             ('input', {'type': 'Submit'}))
 
 
-def search_result(q, by, reverse):
+def search_result(q, by, reverse) -> HTML:
 
-    keys = ['From', 'To', 'Subject', 'Date', 'Size', 'Maildir', 'Msgid']
+    # keys = ['from', 'to', 'subject', 'date', 'size', 'maildir', 'msgid']
+    keys = ['from', 'to', 'subject', 'date']
 
     rows = mu_search(q, by, reverse)
-    body = []
+    body: list[tuple] = []
     for row in rows:
-        rowdata = ['tr']
+        rowdata = []
         for key in keys:
-            rowdata.append(row[key])
-        body.append(rowdata)
+            data = row.get(key, None)
+            if data and key == 'date':
+                dt = datetime.fromtimestamp(int(data))
+                data = dt.strftime('%Y-%m-%d %H:%M')
+            rowdata.append(('td', ('a', {'href': '/?id=' + row['msgid']}, data)))
+        body.append(('tr', rowdata))
+
+
+    # TODO print number of results
 
     return ('table',
             ('thead',
              ('tr',
-              [('th', m) for m in keys])),
+              [('th', m.title()) for m in keys])),
             ('tbody',
              body
              ))
 
 
-def search_page(q, by, username=None):
+def search_page(q, by):
     main_body = [search_field(q)]
 
     if q:
         main_body.append(search_result(q, by, False))
 
-    return render_document(page_base(title='Serach',
-                                     body=main_body,
-                                     username=username))
+    return render_document(page_base(title='Search',
+                                     body=main_body))
 
 
-def index_page(username):
+def index_page():
     ids = [
         'CAEzixGsw-4zJ8_ejK_vDgmcQ9s-MbBc-ho+HL4arV4a+ghOOPg@mail.gmail.com',
         'CA+pcBt-gLb0GtbFOjJ5_7Q_WXtqApVPQ9w-3O7GH=VqCEQat6g@mail.gmail.com',
@@ -240,78 +254,91 @@ def index_page(username):
             ]
 
     return render_document(page_base(title='Mail index',
-                                     body=body,
-                                     username=username
-                                     ))
+                                     body=body))
 
 
 passwords: Passwords = password.Passwords(cast(os.PathLike, 'passwords.json'))
 
 
 app = Flask(__name__)
+login_manager.init_app(app)
 app.secret_key = 'THIS IS A RANDOM STRING'
 
 
+class User:
+    def __init__(self, username: str):
+        self._username = username
+        self._authenticated = False
+
+    # @property
+    def is_authenticated(self):
+        return self._authenticated
+
+    # @property
+    def is_active(self):
+        return True
+
+    # @property
+    def is_anonymous(self):
+        return False
+
+    # @property
+    def get_id(self):
+        return self._username
+
+
+
+
+@login_manager.user_loader
+def load_user(user_id):
+    # return User.get(user_id)
+    return User(user_id)
+
+
 @app.route('/')
 def index():
-    username = session.get('username')
-    if not username:
+    if not current_user.is_authenticated:
         return redirect(url_for('login_page_', returnto=request.path))
     if id := request.args.get('id'):
         print("id =", id)
-        response = response_for(''.join(id).replace(' ', '+'),
-                                username)
+        response = response_for(''.join(id).replace(' ', '+'))
     else:
-        response = index_page(username)
+        response = index_page()
     return response
 
 
 @app.route('/search')
+@login_required
 def search_page_():
-    username = session.get('username')
-    if not username:
-        return redirect(url_for('login_page_', returnto=request.path))
     return search_page(request.args.get('q'),
-                       request.args.get('by'),
-                       username)
+                       request.args.get('by', None))
 
 
 @app.route('/login', methods=['GET'])
 def login_page_():
-    if 'username' not in session:
-        body = login_page(request.args.get('returnto'))
-        return render_document(page_base(title='Login', body=body))
-    else:
-        # TODO do something sensible here
-        pass
+    body = login_page(request.args.get('returnto'))
+    return render_document(page_base(title='Login', body=body))
 
 
 @app.route('/login', methods=['POST'])
 def login_form():
-    logged_in = session.get('username')
-
     resp = redirect(request.args.get('returnto', url_for('index')))
-    if logged_in:
-        flash('Already loged in')
-        return resp
 
     username = request.form['username']
     password = request.form['password']
+    user = User(username)
     if passwords.validate(username, password):
-        session['username'] = username
+        login_user(user)
     else:
         flash('Invalid username or password')
     return resp
 
 
 @app.route('/logout', methods=['POST'])
+@login_required
 def logout_form():
-    if not session.get('username'):
-        flash('Not logged in')
-        return redirect(url_for('index'))
-    session.pop('username', None)
-    resp = redirect(url_for('index'))
-    return resp
+    logout_user()
+    return redirect(url_for('index'))
 
 
 if __name__ == '__main__':
diff --git a/mu.py b/mu.py
index 9337747..1f136df 100644
--- a/mu.py
+++ b/mu.py
@@ -12,6 +12,7 @@ import xml.dom
 
 from typing import (
     Literal,
+    Optional,
     Union,
 )
 
@@ -76,7 +77,7 @@ Sortfield = Union[Literal['cc'],
 
 
 
-def mu_search(query, sortfield: Sortfield = 'subject', reverse: bool = False) -> list[dict[str, str]]:
+def mu_search(query, sortfield: Optional[Sortfield] = 'subject', reverse: bool = False) -> list[dict[str, str]]:
     """
     [Parameters]
         query     - Search query as per mu-find(1).