Add 2FA support

tools/databaseinit.py

@@ -30,6 +30,7 @@ def init_db():
     age TEXT,
     website TEXT,
     blog_theme TEXT,
+    totp CHAR(40),
     Token CHAR(64),
     Lost_password_token CHAR(128),
     invitations INTEGER DEFAULT (20),
@@ -96,14 +97,16 @@ def db_migrate():
     blog_theme_col = False
     updated_col = False
     lost_password_token_col = False
-
+    totp_col = False
+    
     for col in db_columns:
         if "invitations" == col[0]:
             invitations_col = True
         if "Lost_password_token" == col[0]:
             lost_password_token_col = True
-       
-
+        if "totp" == col[0]:
+            totp_col = True
+            
     cursor.execute("""SELECT name FROM PRAGMA_TABLE_INFO('Blog_posts');""")
     db_columns = cursor.fetchall()
     for col in db_columns:
@@ -125,6 +128,12 @@ def db_migrate():
         conn.commit()
         print ("Ajout du champ Lost_password_token dans la table Users")
 
+    if not(totp_col):
+        cursor.execute("""ALTER TABLE Users ADD COLUMN totp CHAR(40);""")
+        conn.commit()
+        print ("Ajout du champ totp dans la table Users")
+
+
         
     if not(blog_theme_col):
         cursor.execute("""ALTER TABLE Blog_posts ADD COLUMN blog_theme TEXT;""")
@@ -136,6 +145,6 @@ def db_migrate():
         conn.commit()
         print ("Ajout du champ updated dans la table BLog")
 
-
+        
 
     conn.close()

tools/utils.py

@@ -3,6 +3,8 @@ import sqlite3
 import os
 import string
 import random
+import time
+import pyotp
 
 app = Flask( 'pywallter' )
 app.config.from_pyfile('config.py')
@@ -129,3 +131,13 @@ def gen_token(token_type):
         case "Lost password":
             sample = ''.join(random.sample(digits + letters, 64))
     return sample
+
+def totp_is_valid(code_key, code):
+    res = True
+    
+    if code_key:
+        mytotp = pyotp.TOTP(code_key)        
+        if not(code == mytotp.now() and res):
+            res = False
+
+    return res