This code handles 2FA (runs in normal mode, but could also run in headless mode). This creates a .selenium folder under the user's home folder, which is where the profile is created and cookies are saved.
So after running this script once and capturing the 2FA code, other scripts could be executed without being prompted for 2FA again.
import sys
from pathlib import Path
from urllib.parse import urlparse, urlunparse, urljoin
from selenium import webdriver
from selenium.webdriver.chrome.options import Options
from selenium.webdriver.chrome.service import Service as ChromeService
from selenium.webdriver.support.ui import WebDriverWait
from selenium.common.exceptions import TimeoutException, NoSuchElementException
from selenium.webdriver.support import expected_conditions as ec
from selenium.webdriver.common.by import By
# location of driver
DRIVER_PATH = '/usr/local/bin/chromedriver'
# selenium profile folder under home folder (saves cookies, etc)
SELENIUM_FOLDER = str(Path.joinpath(Path.home(), ".selenium"))
options = Options()
# set selenium profile folder
options.add_argument(f"user-data-dir={SELENIUM_FOLDER}")
options.add_argument("--window-size=1280,800")
service = ChromeService(executable_path=DRIVER_PATH)
driver = webdriver.Chrome(options=options, service=service)
try:
driver.get("https://test.salesforce.com")
# prompt user for their login credentials to sandbox
user = input("Enter your Salesforce username: ")
pw = input("Enter your Salesforce password: ")
# find and populate login form, log user in
username = driver.find_element(By.ID, "username")
password = driver.find_element(By.ID, "password")
login = driver.find_element(By.ID, "Login")
username.send_keys(user)
password.send_keys(pw)
login.click()
# Handle 2FA if prompted in the browser
try:
smc = driver.find_element(By.ID, "smc") # "smc" is the text input id
verification_code = input("\nEnter the verification code sent via SMS: ")
smc.send_keys(verification_code)
driver.find_element(By.ID, "save").click() # "save" is the submit button id
except NoSuchElementException:
# not prompted for 2FA, so do nothing
pass
# Wait for various redirects before hitting the final home page
try:
WebDriverWait(driver, 20).until(ec.url_contains("/lightning/page/home"))
except TimeoutException as e:
driver.quit()
sys.exit(f"\nERROR: unexpected page reached after login: {driver.current_url}\n")
# logged in, continue with scripted interaction
# Very important! This finally block is needed to quit the driver regardless
# of success or error. If you don't quit the driver, the next selenium run
# will have an error, and you'll have to delete the contents of the .selenium
# folder. Which will result in being prompted again for 2FA...
finally:
print("\nClosing selenium...\n")
driver.quit()
