Merge pull request #75 from ivanhrabcak/feature/2fa

Implement support for 2fa (Closes #74)

edupage_api/__init__.py

@@ -10,7 +10,7 @@
 from edupage_api.custom_request import CustomRequest
 from edupage_api.foreign_timetables import ForeignTimetables, LessonSkeleton
 from edupage_api.grades import EduGrade, Grades, Term
-from edupage_api.login import Login
+from edupage_api.login import Login, TwoFactorLogin
 from edupage_api.lunches import Lunch, Lunches
 from edupage_api.messages import Messages
 from edupage_api.module import EdupageModule
@@ -42,26 +42,36 @@ def __init__(self, request_timeout=5):
         self.session = requests.session()
         self.session.request = functools.partial(self.session.request, timeout=request_timeout)
 
-    def login(self, username: str, password: str, subdomain: str):
+    def login(
+        self, username: str, password: str, subdomain: str
+    ) -> Optional[TwoFactorLogin]:
         """Login while specifying the subdomain to log into.
 
         Args:
             username (str): Your username.
             password (str): Your password.
             subdomain (str): Subdomain of your school (https://{subdomain}.edupage.org).
+
+        Returns:
+            Optional[TwoFactorLogin]: Returns `None` if no second factor was needed to login,
+            or the `TwoFactorLogin` object that is used to complete 2fa.
         """
 
-        Login(self).login(username, password, subdomain)
+        return Login(self).login(username, password, subdomain)
 
-    def login_auto(self, username: str, password: str):
+    def login_auto(self, username: str, password: str) -> Optional[TwoFactorLogin]:
         """Login using https://portal.edupage.org. If this doesn't work, please use `Edupage.login`.
 
         Args:
             username (str): Your username.
             password (str): Your password.
+
+        Returns:
+            Optional[TwoFactorLogin]: Returns `None` if no second factor was needed to login,
+            or the `TwoFactorLogin` object that is used to complete 2fa.
         """
 
-        Login(self).login(username, password)
+        return Login(self).login(username, password)
 
     def get_students(self) -> Optional[list[EduStudent]]:
         """Get list of all students in your class.

edupage_api/exceptions.py

@@ -58,4 +58,7 @@ class UnknownServerError(Exception):
     pass
 
 class NotParentException(Exception):
+    pass
+
+class SecondFactorFailedException(Exception):
     pass

edupage_api/login.py

@@ -1,9 +1,111 @@
 import json
+
 from json import JSONDecodeError
+from typing import Optional
+from dataclasses import dataclass
+from requests import Response
+
+from edupage_api.exceptions import BadCredentialsException, MissingDataException, SecondFactorFailedException, RequestError
+from edupage_api.module import Module, EdupageModule
+
+@dataclass
+class TwoFactorLogin:
+    __authentication_endpoint: str
+    __authentication_token: str
+    __csrf_token: str
+    __edupage: EdupageModule
+
+    __code: Optional[str] = None
+
+    def is_confirmed(self):
+        """Check if the second factor process was finished by confirmation with a device.
+
+        If this function returns true, you can safely use `TwoFactorLogin.finish` to finish the second factor authentication process.
+
+        Returns:
+            bool: True if the second factor was confirmed with a device.
+        """
+
+        request_url = f"https://{self.__edupage.subdomain}.edupage.org/login/twofactor?akcia=checkIfConfirmed"
+        response = self.__edupage.session.post(request_url)
+
+        data = response.json()
+        if data.get("status") == "fail":
+            return False
+        elif data.get("status") != "ok":
+            raise MissingDataException(f"Invalid response from edupage's server!: {str(data)}")
+
+        self.__code = data["data"]
+
+        return True
+
+    def resend_notifications(self):
+        """Resends the confirmation notification to all devices."""
+
+        request_url = f"https://{self.__edupage.subdomain}.edupage.org/login/twofactor?akcia=resendNotifs"
+        response = self.__edupage.session.post(request_url)
+
+        data = response.json()
+        if data.get("status") != "ok":
+            raise RequestError(f"Failed to resend notifications: {str(data)}")
+
+
+    def __finish(self, code: str):
+        request_url = f"https://{self.__edupage.subdomain}.edupage.org/login/edubarLogin.php"
+        parameters = {
+            "csrfauth": self.__csrf_token,
+            "t2fasec": code,
+            "2fNoSave": "y",
+            "2fform": "1",
+            "gu": self.__authentication_endpoint,
+            "au": self.__authentication_token
+        }
+
+        response = self.__edupage.session.post(request_url, parameters)
+
+        if "window.location = gu;" in response.text:
+            cookies = self.__edupage.session.cookies.get_dict(f"{self.__edupage.subdomain}.edupage.org")
+
+            Login(self.__edupage).reload_data(
+                self.__edupage.subdomain, 
+                cookies["PHPSESSID"],
+                self.__edupage.username
+            )
+
+            return
+
+        raise SecondFactorFailedException(f"Second factor failed! (wrong/expired code? expired session?)")
 
-from edupage_api.exceptions import BadCredentialsException
-from edupage_api.module import Module
+    def finish(self):
+        """Finish the second factor authentication process.
+        This function should be used when using a device to confirm the login. If you are using email 2fa codes, please use `TwoFactorLogin.finish_with_code`.
 
+        Notes:
+            - This function can only be used after `TwoFactorLogin.is_confirmed` returned `True`.
+            - This function can raise `SecondFactorFailedException` if there is a big delay from calling `TwoFactorLogin.is_confirmed` (and getting `True` as a result) to calling `TwoFactorLogin.finish`.
+
+        Raises:
+            BadCredentialsException: You didn't call and get the `True` result from `TwoFactorLogin.is_confirmed` before calling this function.
+            SecondFactorFailedException: The delay between calling `TwoFactorLogin.is_confirmed` and `TwoFactorLogin.finish` was too long, or there was another error with the second factor authentication confirmation process.
+        """
+
+        if self.__code is None:
+            raise BadCredentialsException("Not confirmed! (you can only call finish after `TwoFactorLogin.is_confirmed` has returned True)")
+
+        self.__finish(self.__code)
+
+    def finish_with_code(self, code: str):
+        """Finish the second factor authentication process.
+        This function should be used when email 2fa codes are used to confirm the login. If you are using a device to confirm the login, please use `TwoFactorLogin.finish`.
+
+        Args:
+            code (str): The 2fa code from your email or from the mobile app.
+
+        Raises:
+            SecondFactorFailedException: An invalid 2fa code was provided.
+        """
+        self.__finish(code)
+
 
 class Login(Module):
     def __parse_login_data(self, data):
@@ -19,20 +121,26 @@ def __parse_login_data(self, data):
         self.edupage.is_logged_in = True
 
         self.edupage.gsec_hash = data.split('ASC.gsechash="')[1].split('"')[0]
-
-    def login(self, username: str, password: str, subdomain: str = "login1"):
-        """Login to your school's Edupage account.
-        If the subdomain is not specified, the https://login1.edupage.org is used.
-
-        Args:
-            username (str): Your username.
-            password (str): Your password.
-            subdomain (str): Subdomain of your school (https://{subdomain}.edupage.org).
-
-        Raises:
-            BadCredentialsException: Your credentials are invalid.
-        """
-
+
+    def __second_factor(self):
+        request_url = f"https://{self.edupage.subdomain}.edupage.org/login/twofactor?sn=1"
+        two_factor_response = self.edupage.session.get(request_url)
+
+        data = two_factor_response.content.decode()
+
+        csrf_token = data.split("csrfauth\" value=\"")[1].split("\"")[0]
+
+        authentication_token = data.split("au\" value=\"")[1].split("\"")[0]
+        authentication_endpoint = data.split("gu\" value=\"")[1].split("\"")[0]
+
+        return TwoFactorLogin(
+            authentication_endpoint,
+            authentication_token,
+            csrf_token,
+            self.edupage
+        )
+
+    def __login(self, username: str, password: str, subdomain: str) -> Response:
         request_url = f"https://{subdomain}.edupage.org/login/index.php"
 
         response = self.edupage.session.get(request_url)
@@ -49,18 +157,54 @@ def login(self, username: str, password: str, subdomain: str = "login1"):
         request_url = f"https://{subdomain}.edupage.org/login/edubarLogin.php"
 
         response = self.edupage.session.post(request_url, parameters)
-        data = response.content.decode()
 
         if "bad=1" in response.url:
             raise BadCredentialsException()
+
+        return response
+
+    def login(
+        self, username: str, password: str, subdomain: str = "login1"
+    ) -> Optional[TwoFactorLogin]:
+        """Login to your school's Edupage account (optionally with 2 factor authentication).
+
+        If you do not have 2 factor authentication set up, this function will return `None`.
+        The login will still work and succeed.
+
+        See the `Edupage.TwoFactorLogin` documentation or the examples for more details
+        of the 2 factor authentication process.
+
+        Args:
+            username (str): Your username.
+            password (str): Your password.
+            subdomain (str): Subdomain of your school (https://{subdomain}.edupage.org).
+
+        Returns:
+            Optional[TwoFactorLogin]: The object that can be used to complete the second factor
+                (or `None` — if the second factor is not set up)
+
+        Raises:
+            BadCredentialsException: Your credentials are invalid.
+            SecondFactorFailed: The second factor login timed out
+                or there was another problem with the second factor.
+        """
+
+        response = self.__login(username, password, subdomain)
+        data = response.content.decode()
 
         if subdomain == "login1":
             subdomain = data.split("-->")[0].split(" ")[-1]
 
-        self.__parse_login_data(data)
         self.edupage.subdomain = subdomain
         self.edupage.username = username
 
+        if "twofactor" not in response.url:
+            # 2FA not needed
+            self.__parse_login_data(data)
+            return
+
+        return self.__second_factor()
+
     def reload_data(self, subdomain: str, session_id: str, username: str):
         request_url = f"https://{subdomain}.edupage.org/user"
 

examples/2fa.py

@@ -0,0 +1,37 @@
+import time
+from edupage_api import Edupage
+from edupage_api.exceptions import SecondFactorFailedException, BadCredentialsException
+
+edupage = Edupage()
+
+USERNAME = "Username"
+PASSWORD = "Password"
+SUBDOMAIN = "Your school's subdomain"
+
+try:
+    second_factor = edupage.login(USERNAME, PASSWORD, SUBDOMAIN)
+    confirmation_method = input(
+        "Choose confirmation method: 1 -> mobile app, 2 -> code: "
+    )
+
+    if confirmation_method == "1":
+        while not second_factor.is_confirmed():
+            time.sleep(0.5)
+        second_factor.finish()
+
+    elif confirmation_method == "2":
+        code = input("Enter 2FA code (or 'resend' to resend the code): ")
+        while code.lower() == "resend":
+            second_factor.resend_notifications()
+            code = input("Enter 2FA code (or 'resend' to resend the code): ")
+        second_factor.finish_with_code(code)
+
+except BadCredentialsException:
+    print("Wrong username or password")
+except SecondFactorFailedException:
+    print("Second factor failed")
+
+if edupage.is_logged_in:
+    print("Logged in")
+else:
+    print("Login failed")