Merge pull request #96 from Zohair-coder/sso

More robust system

Author: Zohair-coder

docker-compose.yml

@@ -1,4 +1,3 @@
-version: "3"
 services:
   scraper:
     build: .

src/helpers.py

@@ -10,15 +10,16 @@ def send_request(
     method: str = "GET",
     json: dict[str, Any] = {},
     headers: dict[str, str] = {},
+    data: dict[str, Any] = {},
 ) -> Response:
     try:
-        resp = send_request_helper(session, url, method, json, headers)
+        resp = send_request_helper(session, url, method, json, headers, data)
         resp.raise_for_status()
     except Exception as ex:
         retries = 10
         for _ in range(retries):
             try:
-                resp = send_request_helper(session, url, method, json, headers)
+                resp = send_request_helper(session, url, method, json, headers, data)
                 resp.raise_for_status()
             except Exception as inner_ex:
                 if (
@@ -51,12 +52,13 @@ def send_request_helper(
     method: str,
     json: dict[str, Any],
     headers: dict[str, str],
+    data: dict[str, Any],
 ) -> Response:
     timeout = 2
     if method == "GET":
         resp = session.get(url, headers=headers, timeout=timeout)
     elif method == "POST":
-        resp = session.post(url, json=json, headers=headers, timeout=timeout)
+        resp = session.post(url, json=json, data=data, headers=headers, timeout=timeout)
     else:
         raise Exception("Invalid method: {}".format(method))
     return resp

src/login.py

@@ -1,13 +1,16 @@
 from requests import Session
 from bs4 import BeautifulSoup, Tag
 import re
+from typing import Any
 
 import config
 import totp
+from helpers import send_request
 
 
 def login_with_drexel_connect(session: Session) -> Session:
-    response = session.get(config.drexel_connect_base_url)
+    response = send_request(session, config.drexel_connect_base_url, method="GET")
+    assert response.status_code == 200, "Failed to get Drexel Connect login page"
     soup = BeautifulSoup(response.text, "html.parser")
 
     csrf_token = extract_csrf_token(soup)
@@ -21,26 +24,44 @@ def login_with_drexel_connect(session: Session) -> Session:
     }
 
     # this should send the credentials and send the MFA request
-    response = session.post(
-        config.drexel_connect_base_url + form_action_path, data=login_payload
+    response = send_request(
+        session,
+        config.drexel_connect_base_url + form_action_path,
+        data=login_payload,
+        method="POST",
     )
+    assert (
+        response.status_code == 200
+    ), "Failed to send request to Drexel Connect with username and password"
 
     soup = BeautifulSoup(response.text, "html.parser")
     data = parse_initial_mfa_page(soup)
 
-    response = session.post(
-        config.drexel_connect_base_url + data["url"], data=data["form-data"]
+    response = send_request(
+        session,
+        config.drexel_connect_base_url + data["url"],
+        data=data["form-data"],
+        method="POST",
     )
+    assert (
+        response.status_code == 200
+    ), "Failed to request MFA code page from Drexel Connect"
     json_response = response.json()
 
     data = {
         json_response["csrfN"]: json_response["csrfV"],
         "_eventId": json_response["actValue"],
     }
 
-    response = session.post(
-        config.drexel_connect_base_url + json_response["flowExURL"], data=data
+    response = send_request(
+        session,
+        config.drexel_connect_base_url + json_response["flowExURL"],
+        data=data,
+        method="POST",
     )
+    assert (
+        response.status_code == 200
+    ), "Failed to receive MFA code page from Drexel Connect"
     soup = BeautifulSoup(response.text, "html.parser")
 
     parsed_data = parse_final_mfa_page(soup)
@@ -53,9 +74,15 @@ def login_with_drexel_connect(session: Session) -> Session:
         "j_mfaToken": totp_code,
     }
 
-    response = session.post(
-        config.drexel_connect_base_url + parsed_data["url"], data=data
+    response = send_request(
+        session,
+        config.drexel_connect_base_url + parsed_data["url"],
+        data=data,
+        method="POST",
     )
+    assert (
+        response.status_code == 200
+    ), "Failed to send MFA code to Drexel Connect (final step)"
 
     return session
 
@@ -98,7 +125,7 @@ def extract_form_action_path(soup: BeautifulSoup) -> str:
     return form_action_path
 
 
-def parse_initial_mfa_page(soup: BeautifulSoup) -> dict[str, str]:
+def parse_initial_mfa_page(soup: BeautifulSoup) -> dict[str, Any]:
     data = {}
 
     # get the first script tag that isn't empty

src/scrape.py

@@ -3,6 +3,8 @@
 import json
 import os
 from typing import Any
+import traceback
+import time
 
 from helpers import send_request
 from parse import parse_subject_page, parse_crn_page
@@ -15,8 +17,23 @@ def scrape(
 ) -> dict[str, dict[str, Any]]:
     session = Session()
 
-    session = login.login_with_drexel_connect(session)
-    assert "shib_idp_session" in session.cookies, "Failed to log in to Drexel Connect"
+    is_logged_into_drexel_connect = False
+    failiure_count = 0
+    while not is_logged_into_drexel_connect:
+        try:
+            session = login.login_with_drexel_connect(session)
+            if "shib_idp_session" in session.cookies:
+                is_logged_into_drexel_connect = True
+                break
+        except Exception:
+            print("Error logging in to Drexel Connect: ")
+            print(traceback.format_exc())
+            print("Trying again...")
+
+        failiure_count += 1
+        if failiure_count > 5:
+            raise Exception("Failed to log in to Drexel Connect after 5 attempts")
+        time.sleep(1)
 
     data: dict[str, dict[str, Any]] = {}