1 files changed, 74 insertions, 19 deletions
diff --git a/src/jwebmail/read_mails.py b/src/jwebmail/read_mails.py
index 2d2b26c..01252ee 100644
--- a/src/jwebmail/read_mails.py
+++ b/src/jwebmail/read_mails.py
@@ -1,4 +1,5 @@
-import redis
+from datetime import datetime, timedelta
+
 from flask import current_app, g
 from flask_login import UserMixin, current_user
 
@@ -13,6 +14,72 @@ class JWebmailUser(UserMixin):
         self.password = password
 
 
+class RedisTimeoutSession:
+    def __init__(self, username, passwd, timeout, port=6379):
+        import redis
+
+        self.timeout = timeout
+        self.conn = redis.Redis(
+            host="localhost",
+            port=port,
+            decode_responses=True,
+            protocol=3,
+            username=username,
+            password=passwd,
+        )
+
+    def set(self, key, value):
+        self.conn.setex(f"jwm:user:{key}", self.timeout, value)
+
+    def get(self, key):
+        return self.conn.getex(f"jwm:user:{key}", self.timeout)
+
+
+class MysqlTimeoutSession:
+    def __init__(self, username, passwd, timeout, port=3306):
+        import mysql.connector
+
+        self.timeout = timeout
+        self.conn = mysql.connector.connect(
+            host="localhost",
+            port=port,
+            username=username,
+            password=passwd,
+            database="jwebmaildb1",
+        )
+
+    def set(self, key, value):
+        timeout = datetime.now() + timedelta(seconds=self.timeout)
+        cur = self.conn.cursor()
+        cur.execute("DELETE FROM session WHERE user = %s", [key])
+        cur.execute("INSERT INTO session VALUES (%s, %s, %s)", [key, value, timeout])
+        self.conn.commit()
+        cur.close()
+
+    def get(self, key):
+        cur = self.conn.cursor()
+        cur.execute("DELETE FROM session WHERE timeout < NOW()")
+        cur.execute("SELECT password FROM session WHERE user = %s", [key])
+        row = cur.fetchone()
+        self.conn.commit()
+        cur.close()
+        if row is None:
+            return None
+        else:
+            return row[0]
+
+
+def select_timeout_session():
+    session_type = current_app.config["JWEBMAIL"]["READ_MAILS"]["SESSION_TYPE"]
+
+    if session_type == "REDIS":
+        return RedisTimeoutSession
+    elif session_type == "MYSQL":
+        return MysqlTimeoutSession
+    else:
+        raise ValueError(f"unknown session_type {session_type!r}")
+
+
 def build_qma(username, password):
     authenticator = current_app.config["JWEBMAIL"]["READ_MAILS"]["AUTHENTICATOR"]
     backend = current_app.config["JWEBMAIL"]["READ_MAILS"]["BACKEND"]
@@ -30,28 +97,16 @@ def login(username, password):
 
 def add_user(user: JWebmailUser):
     passwd = current_app.config["JWEBMAIL"]["READ_MAILS"]["SESSION_STORE_PASSWD"]
-    r = redis.Redis(
-        host="localhost",
-        port=6379,
-        decode_responses=True,
-        protocol=3,
-        username="jwebmail",
-        password=passwd,
-    )
-    r.setex(f"jwm:user:{user.get_id()}", EXPIRATION_SEC, user.password)
+
+    r = select_timeout_session()("jwebmail", passwd, EXPIRATION_SEC)
+    r.set(user.get_id(), user.password)
 
 
 def load_user(username: str) -> JWebmailUser:
     ss_password = current_app.config["JWEBMAIL"]["READ_MAILS"]["SESSION_STORE_PASSWD"]
-    r = redis.Redis(
-        host="localhost",
-        port=6379,
-        decode_responses=True,
-        protocol=3,
-        username="jwebmail",
-        password=ss_password,
-    )
-    passwd = r.getex(f"jwm:user:{username}", EXPIRATION_SEC)
+
+    r = select_timeout_session()("jwebmail", ss_password, EXPIRATION_SEC)
+    passwd = r.get(username)
     if passwd is None:
         return None
     return JWebmailUser(username, passwd)