From b210677359be03d857d322ab4e91e9cb01a63e5b Mon Sep 17 00:00:00 2001
From: Kimiyuki Onaka <kimiyuki95@gmail.com>
Date: Thu, 31 Jan 2019 03:12:51 +0900
Subject: [PATCH] #200: add --check option to login command

---
 onlinejudge/implementation/command/login.py | 33 +++++++++++++--------
 onlinejudge/implementation/main.py          |  1 +
 onlinejudge/service/atcoder.py              |  7 +++++
 onlinejudge/service/codeforces.py           |  6 ++++
 onlinejudge/service/yukicoder.py            |  7 +++++
 onlinejudge/type.py                         |  3 ++
 6 files changed, 45 insertions(+), 12 deletions(-)

diff --git a/onlinejudge/implementation/command/login.py b/onlinejudge/implementation/command/login.py
index 926d8ad0..ab8d70ad 100644
--- a/onlinejudge/implementation/command/login.py
+++ b/onlinejudge/implementation/command/login.py
@@ -31,16 +31,25 @@ def login(args: 'argparse.Namespace') -> None:
             log.failure('login for %s: invalid option: --method %s', service.get_name(), args.method)
             sys.exit(1)
 
-    # 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
-
-    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')
-
     with utils.with_cookiejar(utils.new_default_session(), path=args.cookie) as sess:
-        service.login(get_credentials, session=sess, **kwargs)  # type: ignore
+
+        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)
+
+        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
+
+            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')
+
+            service.login(get_credentials, session=sess, **kwargs)  # type: ignore
diff --git a/onlinejudge/implementation/main.py b/onlinejudge/implementation/main.py
index ca82cabb..474f1bc6 100644
--- a/onlinejudge/implementation/main.py
+++ b/onlinejudge/implementation/main.py
@@ -85,6 +85,7 @@ def get_parser() -> argparse.ArgumentParser:
     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')
 
     # submit
diff --git a/onlinejudge/service/atcoder.py b/onlinejudge/service/atcoder.py
index 8f016b87..81090a4d 100644
--- a/onlinejudge/service/atcoder.py
+++ b/onlinejudge/service/atcoder.py
@@ -44,6 +44,13 @@ def login(self, get_credentials: onlinejudge.type.CredentialsProvider, session:
         AtCoderService._report_messages(msgs)
         return 'login' not in resp.url  # AtCoder redirects to the top page if success
 
+    def is_logged_in(self, session: Optional[requests.Session] = None) -> bool:
+        session = session or utils.new_default_session()
+        url = 'https://practice.contest.atcoder.jp/login'
+        resp = _request('GET', url, session=session, allow_redirects=False)
+        msgs = AtCoderService._get_messages_from_cookie(resp.cookies)
+        return bool(msgs)
+
     def get_url(self) -> str:
         return 'https://atcoder.jp/'
 
diff --git a/onlinejudge/service/codeforces.py b/onlinejudge/service/codeforces.py
index a20a8708..13f3ea04 100644
--- a/onlinejudge/service/codeforces.py
+++ b/onlinejudge/service/codeforces.py
@@ -44,6 +44,12 @@ def login(self, get_credentials: onlinejudge.type.CredentialsProvider, session:
             log.failure('Invalid handle or password.')
             return False
 
+    def is_logged_in(self, session: Optional[requests.Session] = None) -> bool:
+        session = session or utils.new_default_session()
+        url = 'https://codeforces.com/enter'
+        resp = utils.request('GET', url, session=session, allow_redirects=False)
+        return resp.status_code == 302
+
     def get_url(self) -> str:
         return 'https://codeforces.com/'
 
diff --git a/onlinejudge/service/yukicoder.py b/onlinejudge/service/yukicoder.py
index 6e88efd4..20dc2a05 100644
--- a/onlinejudge/service/yukicoder.py
+++ b/onlinejudge/service/yukicoder.py
@@ -66,6 +66,13 @@ def login_with_twitter(self, get_credentials: onlinejudge.type.CredentialsProvid
         url = 'https://yukicoder.me/auth/twitter'
         raise NotImplementedError
 
+    def is_logged_in(self, session: Optional[requests.Session] = None, method: Optional[str] = None) -> bool:
+        session = session or utils.new_default_session()
+        url = 'https://yukicoder.me/auth/github'
+        resp = utils.request('GET', url, session=session, allow_redirects=False)
+        assert resp.status_code == 302
+        return 'oauth' not in resp.headers['Location']
+
     def get_url(self) -> str:
         return 'https://yukicoder.me/'
 
diff --git a/onlinejudge/type.py b/onlinejudge/type.py
index 9d3e77a9..d7fe56a5 100644
--- a/onlinejudge/type.py
+++ b/onlinejudge/type.py
@@ -11,6 +11,9 @@ class Service(object):
     def login(self, get_credentials: CredentialsProvider, session: Optional['requests.Session'] = None) -> bool:
         raise NotImplementedError
 
+    def is_logged_in(self, session: Optional['requests.Session'] = None) -> bool:
+        raise NotImplementedError
+
     def get_url(self) -> str:
         raise NotImplementedError