From ee57150807f0857f3f37293faea311d4cf39a13b Mon Sep 17 00:00:00 2001
From: Shihaam Abdul Rahman <shihaam@shihaam.me>
Date: Fri, 10 Jan 2025 07:36:06 +0500
Subject: [PATCH] Add support for profile selection

---
 getcookie.py | 48 ++++++++++++++++++++++++++++++++++++++++--------
 1 file changed, 40 insertions(+), 8 deletions(-)

diff --git a/getcookie.py b/getcookie.py
index 10e0dfb..bc01246 100644
--- a/getcookie.py
+++ b/getcookie.py
@@ -20,7 +20,7 @@ class MIBLogin:
 
     def create_new_session(self):
         self.session = requests.Session()
-        
+
     def get_login_page(self):
         response = self.session.get(f"{self.base_url}/auth", headers=self.headers)
         soup = BeautifulSoup(response.text, 'html.parser')
@@ -42,7 +42,7 @@ class MIBLogin:
     def login(self, username, password, retain=1):
         rtag = self.get_login_page()
         self.get_auth_type(rtag, username, retain)
-        
+
         data = {'rTag': rtag, 'pgf01': username, 'pgf02': password, 'retain': retain}
         headers = {**self.headers, 'Content-Type': 'application/x-www-form-urlencoded'}
         self.session.post(
@@ -54,7 +54,7 @@ class MIBLogin:
     def auth_2fa(self, totp_seed):
         self.session.get(f"{self.base_url}/auth2FA", headers=self.headers)
         totp = pyotp.TOTP(totp_seed)
-        
+
         data = {'otpType': 3, 'otp': totp.now()}
         headers = {**self.headers, 'Content-Type': 'application/x-www-form-urlencoded'}
         self.session.post(
@@ -62,7 +62,7 @@ class MIBLogin:
             headers=headers,
             data=urlencode(data)
         )
-        
+
         cookies = {
             cookie.name: cookie.value 
             for cookie in self.session.cookies 
@@ -72,6 +72,35 @@ class MIBLogin:
         self.last_renewal_time = time.time()
         return cookies
 
+    def get_profile_rtag(self):
+        url = f"{self.base_url}/profiles"
+        response = self.session.get(url, headers=self.headers)
+        if response.status_code == 200:
+            soup = BeautifulSoup(response.text, 'html.parser')
+            profile_card = soup.find('div', {'class': 'card profile-card smooth smooth-shadow'})
+            if profile_card:
+                return profile_card.get('data-rt')
+            else:
+                raise ValueError("Failed to fetch profile rTag.")
+        else:
+            raise ValueError(f"Failed to retrieve profiles page. Status code: {response.status_code}")
+
+    def switch_profile(self, profile_id, profile_type):
+        rtag = self.get_profile_rtag()
+        url = f"{self.base_url}/aProfileHandler/switchProfile"
+        data = {'rTag': rtag, 'profileId': profile_id, 'profileType': profile_type}
+        headers = {**self.headers, 'Content-Type': 'application/x-www-form-urlencoded'}
+
+        response = self.session.post(url, headers=headers, data=urlencode(data))
+        if response.status_code == 200:
+            result = response.json()
+            if result.get("success"):
+                print("Profile switched successfully.")
+            else:
+                raise ValueError("Failed to switch profile: " + result.get("reasonText", "Unknown error"))
+        else:
+            raise ValueError(f"Profile switch failed with status code: {response.status_code}")
+
     def clear_cookie(self):
         self.last_cookie = None
         self.create_new_session()
@@ -87,7 +116,10 @@ def renew_cookie_internal():
         mib.clear_cookie()
         mib.login(os.getenv('USERNAME'), os.getenv('PASSWORD'))
         cookies = mib.auth_2fa(os.getenv('TOTP_SEED'))
-        print(f"Cookie renewal completed successfully")
+
+        # Profile selection after successful login and 2FA
+        mib.switch_profile(os.getenv('PROFILE_ID'), os.getenv('PROFILE_TYPE'))
+        print(f"Cookie renewal and profile switch completed successfully")
         return cookies
     except Exception as e:
         print(f"Error during cookie renewal: {str(e)}")
@@ -97,7 +129,7 @@ def renew_cookie_internal():
 def get_cookie():
     client_ip = request.remote_addr
     print(f"Providing cookie to: {client_ip}")
-    
+
     if mib.last_cookie:
         return jsonify(mib.last_cookie)
     else:
@@ -110,12 +142,12 @@ def get_cookie():
 def new_cookie():
     client_ip = request.remote_addr
     current_time = time.time()
-    
+
     # Check if less than 60 seconds since last renewal
     if current_time - mib.last_renewal_time < 60:
         print(f"Providing existing cookie to {client_ip} (last renewal was {int(current_time - mib.last_renewal_time)} seconds ago)")
         return jsonify(mib.last_cookie)
-    
+
     print(f"Renewing cookie as requested by: {client_ip}")
     cookies = renew_cookie_internal()
     if "error" in cookies: