You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
91 lines
2.6 KiB
91 lines
2.6 KiB
2 years ago
|
##
|
||
|
|
## lib/auth.py
|
||
|
|
## Provides functions for Calnet authentication over
|
||
|
|
## Selenium, including cookies export.
|
||
|
|
##
|
||
|
|
## Copyright (c) 2022 Kevin Mo
|
||
|
|
##
|
||
|
|
|
||
|
|
import pickle
|
||
|
|
from selenium.webdriver.common.by import By
|
||
|
|
from selenium.webdriver.support.ui import WebDriverWait
|
||
|
|
from selenium.webdriver.support import expected_conditions as EC
|
||
|
|
|
||
|
|
COOKIES_NAME = "cookies.pkl"
|
||
|
|
|
||
|
|
def save_cookies(driver):
|
||
|
|
"""
|
||
|
|
Save cookies given a Selenium driver.
|
||
|
|
"""
|
||
|
|
pickle.dump(driver.get_cookies(), open(COOKIES_NAME, "wb"))
|
||
|
|
return True
|
||
|
|
|
||
|
|
def load_cookies(driver):
|
||
|
|
"""
|
||
|
|
Load cookies given a Selenium driver.
|
||
|
|
"""
|
||
|
|
try:
|
||
|
|
cookies = pickle.load(open(COOKIES_NAME, "rb"))
|
||
|
|
for cookie in cookies:
|
||
|
|
driver.add_cookie(cookie)
|
||
|
|
except IOError:
|
||
|
|
pass
|
||
|
|
|
||
|
|
return driver
|
||
|
|
|
||
|
|
def check_calnet_auth(driver):
|
||
|
|
"""
|
||
|
|
Returns whether CalNet authentication is ready.
|
||
|
|
"""
|
||
|
|
return "auth.berkeley.edu" in driver.current_url
|
||
|
|
|
||
|
|
def perform_calnet_auth(driver, cid, pwd):
|
||
|
|
"""
|
||
|
|
Perform authentication with CalNet, given
|
||
|
|
that driver is at auth.berkeley.edu.
|
||
|
|
"""
|
||
|
|
print("Performing Calnet authentication")
|
||
|
|
print("Current URL:", driver.current_url)
|
||
|
|
|
||
|
|
if not check_calnet_auth(driver):
|
||
|
|
raise Exception("Cannot perform authentication at current state.")
|
||
|
|
|
||
|
|
# Target user input and then submit
|
||
|
|
cid_box = driver.find_element_by_id("username")
|
||
|
|
pwd_box = driver.find_element_by_id("password")
|
||
|
|
submit_box = driver.find_element_by_id("submit")
|
||
|
|
|
||
|
|
cid_box.send_keys(cid)
|
||
|
|
pwd_box.send_keys(pwd)
|
||
|
|
submit_box.click()
|
||
|
|
|
||
|
|
# Handle incorrect login attempt
|
||
|
|
if "auth.berkeley.edu" in driver.current_url:
|
||
|
|
print("Incorrect login attempt detected.")
|
||
|
|
raise Exception("Incorrect CalNet credentials")
|
||
|
|
|
||
|
|
# Check the presence of Duo 2FA
|
||
|
|
try:
|
||
|
|
if "duosecurity.com" not in driver.current_url:
|
||
|
|
return True
|
||
|
|
|
||
|
|
print()
|
||
|
|
print("IMPORTANT: Complete the 2FA process on the automated browser.")
|
||
|
|
print("Complete the process through push notification, security key, etc.")
|
||
|
|
print()
|
||
|
|
|
||
|
|
# Stall for success or trust box
|
||
|
|
wait = WebDriverWait(driver, 30)
|
||
|
|
trust_box = wait.until(EC.element_to_be_clickable((By.ID, "trust-browser-button")))
|
||
|
|
print("Detected trust browser button, clicking...")
|
||
|
|
trust_box.click()
|
||
|
|
|
||
|
|
wait = WebDriverWait(driver, 10)
|
||
|
|
if not wait.until(EC.url_contains("duosecurity.com")):
|
||
|
|
raise Exception("2FA did not complete successfully")
|
||
|
|
else:
|
||
|
|
print("Detected redirect, authentication is a success!")
|
||
|
|
except:
|
||
|
|
print("Timed out or could not locate Duo 2FA prompt.")
|
||
|
|
|
||
|
|
|