feat: Parse Authorization HTML with BeautifulSoup4

GuyKh · GuyKh · commit 94f98a53ff0d · 2025-06-06T10:06:39.000+03:00
diff --git a/iec_api/commons.py b/iec_api/commons.py
@@ -242,5 +242,7 @@ async def on_request_chunk_sent_debug(
 
 
 async def on_request_end_debug(session: aiohttp.ClientSession, context, params: aiohttp.TraceRequestEndParams):
-    logger.debug(f"HTTP {params.method} call from {params.url} - Response <{params.response.status}>: \
-        {await params.response.text()}")
+    logger.debug(
+        f"HTTP {params.method} call from {params.url} - Response <{params.response.status}>: \
+        {await params.response.text()}"
+    )
diff --git a/iec_api/const.py b/iec_api/const.py
@@ -21,7 +21,7 @@
 
 HEADERS_WITH_AUTH = HEADERS_NO_AUTH.copy()  # Make a copy of the original dictionary
 HEADERS_WITH_AUTH["Authorization"] = "Bearer 1234"
-HEADERS_WITH_AUTH["Cookie"] = "ARRAffinity=?; " "ARRAffinitySameSite=?;" " GCLB=?"
+HEADERS_WITH_AUTH["Cookie"] = "ARRAffinity=?; ARRAffinitySameSite=?; GCLB=?"
 
 TIMEZONE = pytz.timezone("Asia/Jerusalem")
 IEC_API_BASE_URL = "https://iecapi.iec.co.il/api/"
diff --git a/iec_api/login.py b/iec_api/login.py
@@ -3,7 +3,6 @@
 import json
 import logging
 import random
-import re
 import string
 import time
 from typing import Optional, Tuple
@@ -12,6 +11,7 @@
 import jwt
 import pkce
 from aiohttp import ClientSession
+from bs4 import BeautifulSoup
 
 from iec_api import commons
 from iec_api.models.exceptions import IECLoginError
@@ -47,10 +47,20 @@ async def authorize_session(session: ClientSession, session_token) -> str:
     authorize_response = await commons.send_non_json_get_request(
         session=session, url=cmd_url, encoding="unicode-escape"
     )
-    code = re.findall(
-        r"<input type=\"hidden\" name=\"code\" value=\"(.+)\"/>",
-        authorize_response.encode("latin1").decode("utf-8"),
-    )[0]
+
+    # A) Validate that the response is indeed an HTML
+    if not authorize_response.strip().startswith("<!DOCTYPE html>") and not authorize_response.strip().startswith(
+        "<html"
+    ):
+        raise IECLoginError(-1, "Autorize Response is not an HTML document")
+
+    # B) Use BeautifulSoup to extract the code value
+    soup = BeautifulSoup(authorize_response, "html.parser")
+    code_input = soup.find("input", {"name": "code"})
+    if not code_input:
+        raise IECLoginError(-1, "Code input not found in Autorize HTML response")
+
+    code = code_input.get("value")
     return code
 
 
@@ -183,8 +193,7 @@ async def manual_authorization(session: ClientSession, id_number) -> Optional[JW
         raise IECLoginError(-1, "Failed to send OTP, no state_token")
 
     otp_code = await commons.read_user_input("Enter your OTP code: ")
-    code = await authorize_session(session, otp_code)
-    jwt_token = await verify_otp_code(session, factor_id, state_token, code)
+    jwt_token = await verify_otp_code(session, factor_id, state_token, otp_code)
     logger.debug(
         f"Access token: {jwt_token.access_token}\n"
         f"Refresh token: {jwt_token.refresh_token}\n"
diff --git a/iec_api/models/response_descriptor.py b/iec_api/models/response_descriptor.py
@@ -11,7 +11,7 @@ class ResponseDescriptor(DataClassDictMixin):
     """Response Descriptor"""
 
     is_success: bool = field(metadata=field_options(alias="isSuccess"))
-    code: Optional[str]
+    code: Optional[str] = None
     description: Optional[str] = None
 
 
diff --git a/pyproject.toml b/pyproject.toml
@@ -23,6 +23,7 @@ pkce = "^1.0.3"
 aiohttp = "^3.9.1"
 aiofiles = ">=23.2.1,<25.0.0"
 pytz = "^2024.1"
+beautifulsoup4 = "^4.13.4"
 
 [tool.poetry.group.dev.dependencies]
 pytest = "8.3.5"
