diff --git a/main.py b/main.py index 7ceff18..1c7ac8f 100644 --- a/main.py +++ b/main.py @@ -92,8 +92,8 @@ def main(opt): continue # * run command - keep_login_status(driver, login_info) - action_agent.run(action_code) + if keep_login_status(driver, login_info): + action_agent.run(action_code) except KeyboardInterrupt: pass diff --git a/utils/utils.py b/utils/utils.py index fbd0016..e94708b 100644 --- a/utils/utils.py +++ b/utils/utils.py @@ -8,6 +8,7 @@ if sys.version_info >= (3, 11): else: import tomli as tomllib from selenium import webdriver +from selenium.common.exceptions import NoAlertPresentException, UnexpectedAlertPresentException from selenium.webdriver.common.by import By from selenium.webdriver.support.ui import Select @@ -44,7 +45,7 @@ class LoginInfo(): "please enter your passwd: ") -def keep_login_status(driver: webdriver.Chrome, login_info: LoginInfo): +def keep_login_status(driver: webdriver.Chrome, login_info: LoginInfo) -> bool: # * Close all windows except the main window. while len(driver.window_handles) > 1: driver.switch_to.window(driver.window_handles[1]) @@ -77,33 +78,55 @@ def keep_login_status(driver: webdriver.Chrome, login_info: LoginInfo): # * Press the submit button. submit_btn = driver.find_element(By.ID, LOGIN_BUTTON_ID) submit_btn.click() + time.sleep(1) - time.sleep(3) + # * Check if login failed (alert popup) + try: + alert = driver.switch_to.alert + print(f"Login error message: {alert.text}") + alert.accept() + return False # Skip remaining logic, login failed + except NoAlertPresentException: + pass # No alert, proceed + + time.sleep(2) - # * If login has pop a new window, switch main window to the new one. login_page_handle = "" main_page_handle = "" for handle in driver.window_handles: driver.switch_to.window(handle) - if "CHI MotorWeb - " in driver.title: + try: + title = driver.title + except UnexpectedAlertPresentException: + try: + alert = driver.switch_to.alert + print(f"Unexpected alert: {alert.text}") + alert.accept() + continue + except NoAlertPresentException: + continue + + if "CHI MotorWeb - " in title: main_page_handle = handle - elif "CHI MOTOR WEB ERP 登入" in driver.title: + elif "CHI MOTOR WEB ERP 登入" in title: login_page_handle = handle - # * get the page handle that should be stay - if main_page_handle: - stay_page_handle = main_page_handle - elif login_page_handle: - stay_page_handle = login_page_handle - else: - stay_page_handle = top_page + stay_page_handle = main_page_handle or login_page_handle or top_page # * close unnecessary pages for handle in driver.window_handles: if handle != stay_page_handle: - driver.switch_to.window(handle) - driver.close() + try: + driver.switch_to.window(handle) + driver.close() + except Exception as e: + print(f"Error closing window {handle}: {e}") driver.switch_to.window(stay_page_handle) driver.maximize_window() + + return True + + elif "CHI MotorWeb - " in driver.title: + return True