More checks on TFA

V6.py

@@ -391,17 +391,22 @@ def login_part_1():
     )
     # 2FA
     try:
-        wait_until_visible(By.ID, "idTxtBx_SAOTCC_OTC", browser=driver, timeout=5)
+        if not wait_until_visible(By.ID, "idTxtBx_SAOTCC_OTC", browser=driver, timeout=5, raise_error=False):
+            custom_sleep(2)
+            return
+
+        tfa = config.UserCredentials.get_tfa()
+        if tfa is None:
+            error("2FA needed but no code available for this account, sending error")
+            raise ValueError("2FA needed but no code available for this account")
+        else:
+            a2f_code = tfa.now()
 
-        a2f_code = config.UserCredentials.get_tfa().now()
         info(f"Need 2FA, I have code: {a2f_code}")
         send_wait_and_confirm(
             driver.find_element(By.ID, "idTxtBx_SAOTCC_OTC"),
             a2f_code
         )
-    except NoSuchElementException:
-        custom_sleep(2)
-        pass
     except Exception as err:
         log_error(err)
 
@@ -704,7 +709,7 @@ def json_start(json_entry, cred: UserCredentials):
             config.WebDriver.switch_to_driver("PC")
             driver = config.WebDriver.driver
             try:
-                if str(account_id) in json_entry["unban"]:
+                if "unban" in json_entry and str(account_id) in json_entry["unban"]:
                     login_part_1()
                     info("\nGO TO example.com TO PROCEED or wait 1200 secs.")
                     for _ in range(1200):

modules/Classes/UserCredentials.py

@@ -30,7 +30,8 @@ class UserCredentials:
 
     def get_tfa(self):
         if not self.tfa_enable():
-            warning("Warning: TFA is not enabled. Calling get_tfa is an expected behaviour.")
+            warning("Warning: TFA is not enabled. Can't get a TFA code.")
+            return None
         return TOTP(self.data[self.current]["2fa"])
 
     def next_account(self):

modules/driver_tools.py

@@ -71,11 +71,12 @@ def send_wait_and_confirm(element, keys: str) -> None:
 
 
 # Wait for the presence of the element identifier or [timeout]s
-def wait_until_visible(search_by: str, identifier: str, timeout: int = 20, browser=None) -> bool:
+def wait_until_visible(search_by: str, identifier: str, timeout: int = 20, browser=None, raise_error=True) -> bool:
     try:
         WebDriverWait(browser, timeout).until(
             expected_conditions.visibility_of_element_located((search_by, identifier)), "element not found")
         return True
     except TimeoutException as e:
+        if raise_error:
             error(f"element {identifier} not found after {timeout}s")
         return False