diff --git a/pytr/account.py b/pytr/account.py index 547b8ec..49c22ca 100644 --- a/pytr/account.py +++ b/pytr/account.py @@ -2,9 +2,10 @@ import sys from pygments import highlight, lexers, formatters import time +import pathlib from getpass import getpass -from pytr.api import TradeRepublicApi, CREDENTIALS_FILE +from pytr.api import TradeRepublicApi, CREDENTIALS_FILE, COOKIES_FILE from pytr.utils import get_logger @@ -17,26 +18,35 @@ 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, save_credentials=False, credentials_file=None, save_cookies=True, cookies_file=None): ''' If web is true, use web login method as else simulate app login. Check if credentials file exists else create it. If no parameters are set but are needed then ask for input ''' log = get_logger(__name__) - save_cookies = True + credentials_file = pathlib.Path(credentials_file) if credentials_file else CREDENTIALS_FILE - if phone_no is None and CREDENTIALS_FILE.is_file(): + if credentials_file.is_file(): log.info('Found credentials file') - with open(CREDENTIALS_FILE) as f: + with open(credentials_file) as f: lines = f.readlines() - phone_no = lines[0].strip() - pin = lines[1].strip() - phone_no_masked = phone_no[:-8] + '********' - pin_masked = len(pin) * '*' + phone_no_cf = lines[0].strip() + pin_cf = lines[1].strip() + phone_no_masked = phone_no_cf[:-8] + '********' + pin_masked = len(pin_cf) * '*' log.info(f'Phone: {phone_no_masked}, PIN: {pin_masked}') else: - CREDENTIALS_FILE.parent.mkdir(parents=True, exist_ok=True) + phone_no_cf = None + pin_cf = None + + different_account = False + if phone_no is not None and phone_no_cf is not None and phone_no != phone_no_cf: + log.info('Phone number different from credential files. Assuming different account.') + different_account = True + + if phone_no is None and phone_no_cf is None: + credentials_file.parent.mkdir(parents=True, exist_ok=True) if phone_no is None: log.info('Credentials file not found') print('Please enter your TradeRepublic phone number in the format +4912345678:') @@ -48,23 +58,21 @@ 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': - with open(CREDENTIALS_FILE, 'w') as f: - f.writelines([phone_no + '\n', pin + '\n']) - - log.info(f'Saved credentials in {CREDENTIALS_FILE}') - - else: + if save_credentials: + with open(credentials_file, 'w') as f: + f.writelines([phone_no + '\n', pin + '\n']) + log.info(f'Saved credentials in {credentials_file}') + else: + if different_account: save_cookies = False - log.info('Credentials not saved') + log.info('Credentials not saved') - tr = TradeRepublicApi(phone_no=phone_no, pin=pin, save_cookies=save_cookies) + tr = TradeRepublicApi(phone_no=phone_no, pin=pin, save_cookies=save_cookies, + credentials_file=credentials_file, cookies_file=cookies_file) if web: # Use same login as app.traderepublic.com - if tr.resume_websession(): + if not different_account and tr.resume_websession(): log.info('Web session resumed') else: try: diff --git a/pytr/api.py b/pytr/api.py index 20ae661..13c5e77 100644 --- a/pytr/api.py +++ b/pytr/api.py @@ -235,7 +235,12 @@ def resume_websession(self): # Only attempt to load if the cookie file exists. if self._cookies_file.exists(): # Loads session cookies too (expirydate=0). - self._websession.cookies.load(ignore_discard=True, ignore_expires=True) + try: + self._websession.cookies.load(ignore_discard=True, ignore_expires=True) + except: + # cookie file invalid. Delete it. + COOKIES_FILE.unlink() + return False self._weblogin = True try: self.settings() diff --git a/pytr/main.py b/pytr/main.py index b642f22..ec602e0 100644 --- a/pytr/main.py +++ b/pytr/main.py @@ -51,6 +51,9 @@ def formatter(prog): parser_login_args.add_argument('--applogin', help='Use app login instead of web login', action='store_true') parser_login_args.add_argument('-n', '--phone_no', help='TradeRepublic phone number (international format)') parser_login_args.add_argument('-p', '--pin', help='TradeRepublic pin') + parser_login_args.add_argument('-s', '--save_credentials', action='store_true', help='Save credentials') + parser_login_args.add_argument('-c', '--credentials', help='Storage file for credentials') + parser_login_args.add_argument('-k', '--cookies', help='Storage file for cookies') # login info = ( @@ -191,9 +194,12 @@ def main(): log = get_logger(__name__, args.verbosity) log.setLevel(args.verbosity.upper()) log.debug('logging is set to debug') + login_kwargs = dict(phone_no=args.phone_no, pin=args.pin, web=not args.applogin, + save_credentials=args.save_credentials, credentials_file=args.credentials, + cookies_file=args.cookies) if args.command == 'login': - login(phone_no=args.phone_no, pin=args.pin, web=not args.applogin) + login(**login_kwargs) elif args.command == 'dl_docs': if args.last_days == 0: @@ -201,7 +207,7 @@ def main(): else: since_timestamp = (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(**login_kwargs), args.output, args.format, since_timestamp=since_timestamp, @@ -213,11 +219,11 @@ 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(login(**login_kwargs))).get() elif args.command == 'details': - Details(login(phone_no=args.phone_no, pin=args.pin, web=not args.applogin), args.isin).get() + Details(login(login(**login_kwargs)), args.isin).get() elif args.command == 'portfolio': - p = Portfolio(login(phone_no=args.phone_no, pin=args.pin, web=not args.applogin)) + p = Portfolio(login(login(**login_kwargs))) p.get() if args.output is not None: p.portfolio_to_csv(args.output)