From e41470a7fcd1fcb9f539464bdd26c73d58b64ea3 Mon Sep 17 00:00:00 2001
From: Christoph Langer <christoph.langer@sap.com>
Date: Wed, 29 Jan 2025 22:55:19 +0100
Subject: [PATCH] Improve handling of multiple credentials (#144)

---
 pytr/account.py | 12 ++++--------
 pytr/api.py     |  9 ++++++---
 pytr/main.py    | 43 ++++++++++++++++++++++++++++++++++++++-----
 3 files changed, 48 insertions(+), 16 deletions(-)

diff --git a/pytr/account.py b/pytr/account.py
index bb535db..6b104c3 100644
--- a/pytr/account.py
+++ b/pytr/account.py
@@ -19,10 +19,10 @@ def get_settings(tr):
         return formatted_json
 
 
-def login(phone_no=None, pin=None, web=True):
+def login(phone_no=None, pin=None, web=True, store_credentials=False):
     """
-    If web is true, use web login method as else simulate app login.
-    Check if credentials file exists else create it.
+    If web is true, use web login method, else simulate app login.
+    Handle credentials parameters and store to credentials file if requested.
     If no parameters are set but are needed then ask for input
     """
     log = get_logger(__name__)
@@ -52,17 +52,13 @@ def login(phone_no=None, pin=None, web=True):
             print("Please enter your TradeRepublic pin:")
             pin = getpass(prompt="Pin (Input is hidden):")
 
-        print('Save credentials? Type "y" to save credentials:')
-        save = input()
-        if save == "y":
+        if store_credentials:
             with open(CREDENTIALS_FILE, "w") as f:
                 f.writelines([phone_no + "\n", pin + "\n"])
 
             log.info(f"Saved credentials in {CREDENTIALS_FILE}")
-
         else:
             save_cookies = False
-            log.info("Credentials not saved")
 
     tr = TradeRepublicApi(phone_no=phone_no, pin=pin, save_cookies=save_cookies)
 
diff --git a/pytr/api.py b/pytr/api.py
index ae471c7..51ce9c7 100644
--- a/pytr/api.py
+++ b/pytr/api.py
@@ -99,9 +99,6 @@ def __init__(
         self._credentials_file = (
             pathlib.Path(credentials_file) if credentials_file else CREDENTIALS_FILE
         )
-        self._cookies_file = (
-            pathlib.Path(cookies_file) if cookies_file else COOKIES_FILE
-        )
 
         if not (phone_no and pin):
             try:
@@ -117,6 +114,12 @@ def __init__(
             self.phone_no = phone_no
             self.pin = pin
 
+        self._cookies_file = (
+            pathlib.Path(cookies_file)
+            if cookies_file
+            else BASE_DIR / f"cookies.{self.phone_no}.txt"
+        )
+
         self.keyfile = keyfile if keyfile else KEY_FILE
         try:
             with open(self.keyfile, "rb") as f:
diff --git a/pytr/main.py b/pytr/main.py
index 8d80eda..499b687 100644
--- a/pytr/main.py
+++ b/pytr/main.py
@@ -65,6 +65,12 @@ def formatter(prog):
         "-n", "--phone_no", help="TradeRepublic phone number (international format)"
     )
     parser_login_args.add_argument("-p", "--pin", help="TradeRepublic pin")
+    parser_login_args.add_argument(
+        "--store_credentials",
+        help="Store credentials (Phone number, pin, cookies) for next usage",
+        action="store_true",
+        default=False,
+    )
 
     # sort
     parser_sort_export = argparse.ArgumentParser(add_help=False)
@@ -245,7 +251,12 @@ def main():
     log.debug("logging is set to debug")
 
     if args.command == "login":
-        login(phone_no=args.phone_no, pin=args.pin, web=not args.applogin)
+        login(
+            phone_no=args.phone_no,
+            pin=args.pin,
+            web=not args.applogin,
+            store_credentials=args.store_credentials,
+        )
 
     elif args.command == "dl_docs":
         if args.last_days == 0:
@@ -255,7 +266,12 @@ def main():
                 datetime.now().astimezone() - timedelta(days=args.last_days)
             ).timestamp()
         dl = DL(
-            login(phone_no=args.phone_no, pin=args.pin, web=not args.applogin),
+            login(
+                phone_no=args.phone_no,
+                pin=args.pin,
+                web=not args.applogin,
+                store_credentials=args.store_credentials,
+            ),
             args.output,
             args.format,
             since_timestamp=since_timestamp,
@@ -268,15 +284,32 @@ def main():
         # TODO
         print("Not implemented yet")
     elif args.command == "get_price_alarms":
-        Alarms(login(phone_no=args.phone_no, pin=args.pin, web=not args.applogin)).get()
+        Alarms(
+            login(
+                phone_no=args.phone_no,
+                pin=args.pin,
+                web=not args.applogin,
+                store_credentials=args.store_credentials,
+            )
+        ).get()
     elif args.command == "details":
         Details(
-            login(phone_no=args.phone_no, pin=args.pin, web=not args.applogin),
+            login(
+                phone_no=args.phone_no,
+                pin=args.pin,
+                web=not args.applogin,
+                store_credentials=args.store_credentials,
+            ),
             args.isin,
         ).get()
     elif args.command == "portfolio":
         p = Portfolio(
-            login(phone_no=args.phone_no, pin=args.pin, web=not args.applogin)
+            login(
+                phone_no=args.phone_no,
+                pin=args.pin,
+                web=not args.applogin,
+                store_credentials=args.store_credentials,
+            )
         )
         p.get()
         if args.output is not None: