From 94563bfefcdfe371f424d3f8b0698b0982c11246 Mon Sep 17 00:00:00 2001
From: aitzol <aitzol@lainoa.eus>
Date: Sat, 18 Nov 2023 23:03:40 +0100
Subject: [PATCH] 2fa-0.7

---
 app.py         |  7 +++++--
 libs/helper.py | 23 +++++++++++++++++++----
 2 files changed, 24 insertions(+), 6 deletions(-)

diff --git a/app.py b/app.py
index f93756c..cad796b 100644
--- a/app.py
+++ b/app.py
@@ -296,7 +296,7 @@ def post_enable_2fa():
         if(not newSession().get()['secureAuth']):
             try:
                 username=newSession().get()['username']
-                add_auth_attribute_step1(username, tools.generate_secret(), action='enable')
+                add_auth_attribute_step1(username, tools.gen_secret(), action='enable')
             except Error as e:
                 #add_auth_attribute_step1(newSession().get()['username'], None, None)
                 reload(newSession().get()['username'], None, None)
@@ -478,7 +478,10 @@ def login_user_ldap(conf, username, password):
         if is_trusted_device(conf, user_dn):
             newSession().set(get_user_data(user_dn, c))
             #update timestamp + ip address
-            update_login_info(conf, user_dn)        
+            update_login_info(conf, user_dn)
+            #check if exists 2fa qr image
+            if(newSession().get()['secureAuth']):
+                tools.gen_qr(newSession().get()['authCode'])
             LOG.debug("%s logged in to %s" % (username, conf['base']))
 
 #LOGOUT
diff --git a/libs/helper.py b/libs/helper.py
index d3220c6..e8b9214 100644
--- a/libs/helper.py
+++ b/libs/helper.py
@@ -5,6 +5,7 @@ import re
 from onetimepass import valid_totp
 from secrets import choice
 import segno
+from os import path
 
 class Tools():
     
@@ -47,12 +48,26 @@ class Tools():
         return(bool(re.fullmatch(regex, e)))
     
     # 2FA
-    def generate_secret(self):  # Function to return a random string with length 16.
+    def gen_qr(self, secret):
+        if(not path.isfile('static/tmp/'+secret+'.png')):
+            qrcode = segno.make(secret, micro=False)
+            qrcode.save('static/tmp/'+secret+'.png', scale=10)
+
+    def gen_secret(self):  # Function to return a random string with length 16.
         secret = ''
         while len(secret) < 16:
             secret += choice('ABCDEFGHIJKLMNOPQRSTUVWXYZ234567')
-        qrcode = segno.make(secret, micro=False)
-        qrcode.save('static/tmp/'+secret+'.png', scale=10)
+        self.gen_qr(secret)
         return secret
 
-tools = Tools()
+    def 2fa_validation(self, otp):
+        authenticated = valid_totp(otp, secret)
+        if authenticated:
+            print('Correct otp, Authenticated!')
+            return True
+        elif not authenticated:
+            print('Wrong otp, please try again.')
+            return False
+
+
+Tools = Tools()