Automate vsb authentication for scraper (#665)

Author: 39bytes

Cargo.lock

@@ -43,6 +43,7 @@ dev: services typeshare
 dev-deps:
   cargo install present
   cargo install typeshare-cli
+  brew install --cask chromedriver
   curl -LsSf https://astral.sh/uv/install.sh | sh
 
 e2e:
@@ -71,13 +72,12 @@ initialize *args: restart-services
 lint *args:
   pnpm run lint {{args}}
 
-load cookie:
-  cargo run --manifest-path crates/scraper/Cargo.toml -- --source=seed \
+load:
+  cargo run --manifest-path tools/scraper/Cargo.toml -- --source=seed \
     --batch-size=5 \
     --scrape-vsb \
     --user-agent "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/111.0.0.0 Safari/537.36" \
-    --course-delay 1000 \
-    --cookie '{{cookie}}'
+    --course-delay 1000
 
 readme:
   present --in-place README.md

justfile

@@ -17,6 +17,9 @@ reqwest = { version = "0.12.15", default-features = false, features = [ "blockin
 scraper = "0.23.1"
 serde = "1.0.219"
 serde_json = "1.0.140"
+thirtyfour = "0.35.0"
+tokio = "1.45.1"
+totp-rs = "5.7.0"
 
 [dev-dependencies]
 include_dir = "0.7.4"

tools/scraper/Cargo.toml

@@ -0,0 +1,129 @@
+use super::*;
+
+const CHROMEDRIVER_PORT: usize = 9515;
+const MAX_ELEM_RETRIES: usize = 5;
+const VSB_LOGIN_URL: &str = "https://vsb.mcgill.ca/login.jsp";
+
+struct Driver(Child);
+
+impl Drop for Driver {
+  fn drop(&mut self) {
+    let _ = self.0.kill();
+    let _ = self.0.wait();
+  }
+}
+
+pub(crate) fn authenticate() -> Result<String> {
+  let email =
+    env::var("VSB_EMAIL").expect("VSB_EMAIL must be specified for scraping");
+
+  let password = env::var("VSB_PASSWORD")
+    .expect("VSB_PASSWORD must be specified for scraping");
+
+  let otp_secret = env::var("VSB_OTP_SECRET")
+    .expect("VSB_OTP_SECRET must be specified for scraping");
+
+  info!("Starting chromedriver server");
+
+  let _chromedriver = Driver(
+    Command::new("chromedriver")
+      .args([format!("--port={}", CHROMEDRIVER_PORT)])
+      .spawn()?,
+  );
+
+  thread::sleep(Duration::from_secs(2));
+
+  info!("Retrieving cookie for VSB authentication...");
+
+  let rt = tokio::runtime::Builder::new_current_thread()
+    .enable_all()
+    .build()?;
+
+  rt.block_on(get_vsb_cookie(email, password, otp_secret))
+}
+
+fn format_cookie(cookies: Vec<Cookie>) -> String {
+  cookies
+    .into_iter()
+    .map(|c| format!("{}={}", c.name, c.value))
+    .collect::<Vec<String>>()
+    .join("; ")
+}
+
+async fn get_vsb_cookie(
+  email: String,
+  password: String,
+  otp_secret: String,
+) -> Result<String> {
+  let mut caps = DesiredCapabilities::chrome();
+  caps.set_headless()?;
+
+  let driver =
+    WebDriver::new(format!("http://localhost:{}", CHROMEDRIVER_PORT), caps)
+      .await?;
+
+  // Need to use `new_unchecked` because Microsoft auth secret length is too short.
+  let totp = TOTP::new_unchecked(
+    totp_rs::Algorithm::SHA1,
+    6,
+    1,
+    30,
+    Secret::Encoded(otp_secret.to_uppercase()).to_bytes()?,
+  );
+
+  driver.goto(VSB_LOGIN_URL).await?;
+
+  info!("Entering email...");
+  let email_field = retry_until_visible(&driver, By::Name("loginfmt")).await?;
+  email_field.send_keys(email).await?;
+  email_field.send_keys(Key::Return).await?;
+  tokio::time::sleep(Duration::from_secs(1)).await;
+
+  info!("Entering password...");
+  let password_field = retry_until_visible(&driver, By::Name("passwd")).await?;
+  password_field.send_keys(password).await?;
+  tokio::time::sleep(Duration::from_secs(1)).await;
+
+  info!("Signing in...");
+  let submit_button =
+    retry_until_visible(&driver, By::Id("idSIButton9")).await?;
+  submit_button.click().await?;
+  tokio::time::sleep(Duration::from_secs(1)).await;
+
+  info!("Entering OTP...");
+  let otp_field = retry_until_visible(&driver, By::Name("otc")).await?;
+  otp_field.send_keys(totp.generate_current()?).await?;
+  otp_field.send_keys(Key::Return).await?;
+  tokio::time::sleep(Duration::from_secs(3)).await;
+
+  info!("Finishing up...");
+  let no_button = retry_until_visible(&driver, By::Id("idBtn_Back")).await?;
+  no_button.click().await?;
+  tokio::time::sleep(Duration::from_secs(1)).await;
+
+  let cookies = driver.get_all_cookies().await?;
+
+  driver.quit().await?;
+
+  Ok(format_cookie(cookies))
+}
+
+async fn retry_until_visible(driver: &WebDriver, by: By) -> Result<WebElement> {
+  let mut retries = 0;
+
+  let mut elem = driver.find(by.clone()).await;
+
+  while elem.is_err() {
+    if retries > MAX_ELEM_RETRIES {
+      break;
+    }
+
+    tokio::time::sleep(Duration::from_secs(1)).await;
+
+    retries += 1;
+
+    elem = driver.find(by.clone()).await;
+  }
+
+  Ok(elem?)
+}