Merge pull request #508 from kmyk/issue/306

#306: support Selenium

onlinejudge/_implementation/command/login.py

@@ -1,8 +1,13 @@
 # Python Version: 3.x
+import datetime
 import getpass
+import http.cookies
 import sys
+import time
 from typing import *
 
+import requests
+
 import onlinejudge
 import onlinejudge._implementation.logging as log
 import onlinejudge._implementation.utils as utils
@@ -11,48 +16,113 @@
     import argparse
 
 
+def login_with_password(service: onlinejudge.type.Service, *, username: Optional[str], password: Optional[str], session: requests.Session) -> None:
+    def get_credentials() -> Tuple[str, str]:
+        nonlocal username, password
+        if username is None:
+            username = input('Username: ')
+        if password is None:
+            password = getpass.getpass()
+        return username, password
+
+    service.login(get_credentials=get_credentials, session=session)
+
+
+# a wrapper class, because selenium.common.exceptions.* is not always imported
+class WebDriverException(Exception):
+    pass
+
+
+def login_with_browser(service: onlinejudge.type.Service, *, session: requests.Session) -> None:
+    try:
+        import selenium.webdriver
+    except ImportError:
+        raise
+
+    try:
+        profile = selenium.webdriver.FirefoxProfile()
+        profile.set_preference("general.useragent.override", session.headers['User-Agent'])
+        driver = selenium.webdriver.Firefox(firefox_profile=profile)
+    except selenium.common.exceptions.WebDriverException as e:
+        raise WebDriverException(e)
+
+    # get cookies via Selenium
+    url = service.get_url_of_login_page()
+    log.info('open with WebDriver: %s', url)
+    driver.get(url)
+    cookies = []  # type: List[Dict[str, str]]
+    try:
+        while driver.current_url:
+            cookies = driver.get_cookies()
+            time.sleep(0.1)
+    except selenium.common.exceptions.NoSuchWindowException:
+        pass  # the window is closed
+
+    # set cookies to the requests.Session
+    log.info('copy cookies from WebDriver')
+    for c in cookies:
+        log.status('set cookie: %s', c['name'])
+        morsel = http.cookies.Morsel()  # type: http.cookies.Morsel
+        morsel.set(c['name'], c['value'], c['value'])
+        morsel.update({key: value for key, value in c.items() if morsel.isReservedKey(key)})
+        if not morsel['expires']:
+            expires = datetime.datetime.now(datetime.timezone.utc).astimezone() + datetime.timedelta(days=180)
+            morsel.update({'expires': expires.strftime('%a, %d-%b-%Y %H:%M:%S GMT')})  # RFC2109 format
+        cookie = requests.cookies.morsel_to_cookie(morsel)
+        session.cookies.set_cookie(cookie)  # type: ignore
+
+
+def is_logged_in_with_message(service: onlinejudge.type.Service, *, session: requests.Session) -> bool:
+    if service.is_logged_in(session=session):
+        log.info('You have already signed in.')
+        return True
+    else:
+        log.warning('You are not signed in.')
+        return False
+
+
 def login(args: 'argparse.Namespace') -> None:
-    # get service
     service = onlinejudge.dispatch.service_from_url(args.url)
     if service is None:
         sys.exit(1)
 
-    # configure
-    kwargs = {}
-    if isinstance(service, onlinejudge.service.yukicoder.YukicoderService):
-        if not args.method:
-            args.method = 'github'
-        if args.method not in ['github', 'twitter']:
-            log.failure('login for yukicoder: invalid option: --method %s', args.method)
-            sys.exit(1)
-        kwargs['method'] = args.method
-    else:
-        if args.method:
-            log.failure('login for %s: invalid option: --method %s', service.get_name(), args.method)
-            sys.exit(1)
-
-    with utils.with_cookiejar(utils.new_session_with_our_user_agent(), path=args.cookie) as sess:
-
-        if args.check:
-            if service.is_logged_in(session=sess):
-                log.info('You have already signed in.')
-            else:
-                log.info('You are not signed in.')
-                sys.exit(1)
+    with utils.with_cookiejar(utils.new_session_with_our_user_agent(), path=args.cookie) as session:
 
+        if is_logged_in_with_message(service, session=session):
+            return
         else:
-            # login
-            def get_credentials() -> Tuple[str, str]:
-                if args.username is None:
-                    args.username = input('Username: ')
-                if args.password is None:
-                    args.password = getpass.getpass()
-                return args.username, args.password
+            if args.check:
+                sys.exit(1)
 
-            log.warning('If you don\'t want to give your password to this program, you can give only your session tokens.')
-            log.info('see: https://github.com/kmyk/online-judge-tools/blob/master/LOGIN_WITH_COOKIES.md')
+        if args.use_browser in ('always', 'auto'):
+            try:
+                login_with_browser(service, session=session)
+            except ImportError:
+                log.error('Selenium is not installed: try $ pip3 install selenium')
+                pass
+            except WebDriverException as e:
+                log.error('%s', e)
+                pass
+            else:
+                if is_logged_in_with_message(service, session=session):
+                    return
+                else:
+                    sys.exit(1)
 
+        if args.use_browser in ('never', 'auto'):
+            if args.use_browser == 'auto':
+                log.warning('use CUI login since Selenium fails')
             try:
-                service.login(get_credentials, session=sess, **kwargs)  # type: ignore
-            except onlinejudge.type.LoginError:
+                login_with_password(service, username=args.username, password=args.password, session=session)
+            except NotImplementedError as e:
+                log.error('%s', e)
                 pass
+            except onlinejudge.type.LoginError:
+                sys.exit(1)
+            else:
+                if is_logged_in_with_message(service, session=session):
+                    return
+                else:
+                    sys.exit(1)
+
+        sys.exit(1)

onlinejudge/_implementation/main.py

@@ -81,16 +81,12 @@ def get_parser() -> argparse.ArgumentParser:
   Topcoder
   HackerRank
   Toph
-
-strings for --method:
-  github                for yukicoder, login via github (default)
-  twitter               for yukicoder, login via twitter (not implementated yet)
 ''')
     subparser.add_argument('url')
     subparser.add_argument('-u', '--username')
     subparser.add_argument('-p', '--password')
     subparser.add_argument('--check', action='store_true', help='check whether you are logged in or not')
-    subparser.add_argument('--method')
+    subparser.add_argument('--use-browser', choices=('always', 'auto', 'never'), default='auto', help='specify whether it uses a GUI web browser to login or not  (default: auto)')
 
     # submit
     subparser = subparsers.add_parser('submit', aliases=['s'], help='submit your solution', formatter_class=argparse.RawTextHelpFormatter, epilog='''\

onlinejudge/service/atcoder.py

@@ -89,6 +89,9 @@ def login(self, *, get_credentials: onlinejudge.type.CredentialsProvider, sessio
             log.failure('Username or Password is incorrect.')
             raise LoginError
 
+    def get_url_of_login_page(self) -> str:
+        return 'https://atcoder.jp/login'
+
     def is_logged_in(self, *, session: Optional[requests.Session] = None) -> bool:
         session = session or utils.get_default_session()
         url = 'https://atcoder.jp/contests/practice/submit'

onlinejudge/service/codeforces.py

@@ -52,6 +52,9 @@ def login(self, *, get_credentials: onlinejudge.type.CredentialsProvider, sessio
             log.failure('Invalid handle or password.')
             raise LoginError('Invalid handle or password.')
 
+    def get_url_of_login_page(self) -> str:
+        return 'https://codeforces.com/enter'
+
     def is_logged_in(self, *, session: Optional[requests.Session] = None) -> bool:
         session = session or utils.get_default_session()
         url = 'https://codeforces.com/enter'

onlinejudge/service/hackerrank.py

@@ -20,40 +20,8 @@
 
 
 class HackerRankService(onlinejudge.type.Service):
-    def login(self, *, get_credentials: onlinejudge.type.CredentialsProvider, session: Optional[requests.Session] = None) -> None:
-        """
-        :raises LoginError:
-        """
-
-        session = session or utils.get_default_session()
-        url = 'https://www.hackerrank.com/auth/login'
-        # get
-        resp = utils.request('GET', url, session=session)
-        if resp.url != url:
-            log.info('You have already signed in.')
-            return
-        # parse
-        soup = bs4.BeautifulSoup(resp.content.decode(resp.encoding), utils.html_parser)
-        csrftoken = soup.find('meta', attrs={'name': 'csrf-token'}).attrs['content']
-        tag = soup.find('input', attrs={'name': 'username'})
-        while tag.name != 'form':
-            tag = tag.parent
-        form = tag
-        # post
-        username, password = get_credentials()
-        form = utils.FormSender(form, url=resp.url)
-        form.set('login', username)
-        form.set('password', password)
-        form.set('remember_me', 'true')
-        form.set('fallback', 'true')
-        resp = form.request(session, method='POST', action='/rest/auth/login', headers={'X-CSRF-Token': csrftoken})
-        resp.raise_for_status()
-        # result
-        if '/auth' not in resp.url:
-            log.success('You signed in.')
-        else:
-            log.failure('You failed to sign in. Wrong user ID or password.')
-            raise LoginError('You failed to sign in. Wrong user ID or password.')
+    def get_url_of_login_page(self) -> str:
+        return 'https://www.hackerrank.com/auth/login'
 
     def is_logged_in(self, *, session: Optional[requests.Session] = None) -> bool:
         session = session or utils.get_default_session()

onlinejudge/service/topcoder.py

@@ -47,6 +47,9 @@ def login(self, *, get_credentials: onlinejudge.type.CredentialsProvider, sessio
             log.failure('Failure')
             raise LoginError
 
+    def get_url_of_login_page(self):
+        return 'https://accounts.topcoder.com/member'
+
     def is_logged_in(self, *, session: Optional[requests.Session] = None) -> bool:
         """
         .. versionadded:: 6.2.0

onlinejudge/service/toph.py

@@ -20,36 +20,8 @@
 
 
 class TophService(onlinejudge.type.Service):
-    def login(self, *, get_credentials: onlinejudge.type.CredentialsProvider, session: Optional[requests.Session] = None) -> None:
-        """
-        :raises LoginError:
-        """
-        session = session or utils.get_default_session()
-        url = 'https://toph.co/login'
-        # get
-        resp = utils.request('GET', url, session=session)
-        if resp.url != url:  # redirected
-            log.info('You are already logged in.')
-            return
-        # parse
-        soup = bs4.BeautifulSoup(resp.content.decode(resp.encoding), utils.html_parser)
-        form = soup.find('form', class_='login-form')
-        log.debug('form: %s', str(form))
-        username, password = get_credentials()
-        form['action'] = '/login'  # to avoid KeyError inside form.request method as Toph does not have any defined action
-        form = utils.FormSender(form, url=resp.url)
-        form.set('handle', username)
-        form.set('password', password)
-        # post
-        resp = form.request(session)
-        resp.raise_for_status()
-
-        resp = utils.request('GET', url, session=session)  # Toph's Location header is not getting the expected value
-        if resp.url != url:
-            log.success('Welcome, %s.', username)
-        else:
-            log.failure('Invalid handle/email or password.')
-            raise LoginError('Invalid handle/email or password.')
+    def get_url_of_login_page(self) -> str:
+        return 'https://toph.co/login'
 
     def is_logged_in(self, *, session: Optional[requests.Session] = None) -> bool:
         session = session or utils.get_default_session()