new login flow

floyd/cli/auth.py

@@ -5,70 +5,97 @@
 import floyd
 from floyd.client.auth import AuthClient
 from floyd.manager.auth_config import AuthConfigManager
+from floyd.manager.login import (
+    has_browser,
+    wait_for_apikey,
+)
 from floyd.model.access_token import AccessToken
 from floyd.model.credentials import Credentials
 from floyd.log import logger as floyd_logger
 
 
-@click.command()
-@click.option('--token', is_flag=True, default=False, help='Just enter token')
-@click.option('--apikey', '-k', default=None, help='Api Key')
-@click.option('--username', '-u', default=None, help='FloydHub username')
-@click.option('--password', '-p', default=None, help='FloydHub password')
-def login(token, apikey, username, password):
-    """
-    Login to FloydHub.
-    """
-    if apikey:
-        user = AuthClient().get_user(apikey, True)
-        AuthConfigManager.set_apikey(username=user.username, apikey=apikey)
-        floyd_logger.info("Login Successful as %s", user.username)
-        return
+def get_access_code_from_password(username, password):
+    # Use username / password login
+    floyd_logger.info("Login with your FloydHub username and password to run jobs.")
+
+    if not username:
+        username = click.prompt('Username', type=str, default=getpass.getuser())
+        username = username.strip()
+        if not username:
+            floyd_logger.info('You entered an empty string. Please make sure you enter your username correctly.')
+            sys.exit(1)
+
+    if not password:
+        password = click.prompt('Password', type=str, hide_input=True)
+        password = password.strip()
+        if not password:
+            floyd_logger.info('You entered an empty string. Please make sure you enter your password correctly.')
+            sys.exit(1)
 
-    elif token:
+    login_credentials = Credentials(username=username,
+                                    password=password)
+    access_token = AuthClient().login(login_credentials)
+    if not access_token:
+        floyd_logger.info("Failed to login")
+        sys.exit(1)
+
+    return access_token
+
+
+def manual_login_success(token, username, password):
+    if token:
         # Login using authentication token
         floyd_logger.info(
-            "Please paste the authentication token from {}/settings/security.".format(floyd.floyd_web_host))
-        access_code = click.prompt('This is an invisible field. Paste token and press ENTER', type=str, hide_input=True)
-        access_code = access_code.strip()
+            "Please paste the authentication token from %s/settings/security.",
+            floyd.floyd_web_host)
+        access_token = click.prompt('This is an invisible field. Paste token and press ENTER', type=str, hide_input=True)
+        access_token = access_token.strip()
 
-        if not access_code:
+        if not access_token:
             floyd_logger.info("Empty token received. Make sure your shell is handling the token appropriately.")
             floyd_logger.info("See docs for help: http://docs.floydhub.com/faqs/authentication/")
             return
 
-        access_code = access_code.strip(" ")
-
+    elif username or password:
+        access_token = get_access_code_from_password(username, password)
     else:
-        # Use username / password login
-        floyd_logger.info("Login with your FloydHub username and password to run jobs.")
+        return False
 
-        if not username:
-            username = click.prompt('Username', type=str, default=getpass.getuser())
-            username = username.strip()
-            if not username:
-                floyd_logger.info('You entered an empty string. Please make sure you enter your username correctly.')
-                sys.exit(1)
+    user = AuthClient().get_user(access_token)
+    AuthConfigManager.set_access_token(
+        AccessToken(username=user.username,
+                    token=access_token))
+    floyd_logger.info("Login Successful as %s", user.username)
+    return True
 
-        if not password:
-            password = click.prompt('Password', type=str, hide_input=True)
-            password = password.strip()
-            if not password:
-                floyd_logger.info('You entered an empty string. Please make sure you enter your password correctly.')
-                sys.exit(1)
-
-        login_credentials = Credentials(username=username,
-                                        password=password)
-        access_code = AuthClient().login(login_credentials)
-        if not access_code:
-            floyd_logger.info("Failed to login")
-            return
 
-    user = AuthClient().get_user(access_code)
-    access_token = AccessToken(username=user.username,
-                               token=access_code)
-    AuthConfigManager.set_access_token(access_token)
-    floyd_logger.info("Login Successful as %s", user.username)
+@click.command()
+@click.option('--token', is_flag=True, default=False, help='Just enter token')
+@click.option('--apikey', '-k', help='Api Key')
+@click.option('--username', '-u', help='FloydHub username')
+@click.option('--password', '-p', help='FloydHub password')
+def login(token, apikey, username, password):
+    """
+    Login to FloydHub.
+    """
+    if manual_login_success(token, username, password):
+        return
+
+    if not apikey:
+        if has_browser():
+            apikey = wait_for_apikey()
+        else:
+            floyd_logger.error(
+                "No browser found, please login manually by creating login key at %s/settings/apikey.",
+                floyd.floyd_web_host)
+            sys.exit(1)
+
+    if apikey:
+        user = AuthClient().get_user(apikey, is_apikey=True)
+        AuthConfigManager.set_apikey(username=user.username, apikey=apikey)
+        floyd_logger.info("Login Successful as %s", user.username)
+    else:
+        floyd_logger.error("Login failed, please see --help for other login options.")
 
 
 @click.command()

floyd/manager/login.py

@@ -0,0 +1,114 @@
+import sys
+import threading
+import socket
+import contextlib
+import webbrowser
+import subprocess
+
+import floyd
+from floyd.log import logger as floyd_logger
+
+if sys.version_info[0] < 3:
+    import urlparse
+    from BaseHTTPServer import (
+        BaseHTTPRequestHandler,
+        HTTPServer,
+    )
+    from Queue import Queue, Empty as QueueEmpty
+else:
+    from urllib import parse as urlparse
+    from http.server import (
+        BaseHTTPRequestHandler,
+        HTTPServer,
+    )
+    from queue import Queue, Empty as QueueEmpty
+
+
+class LoginServer(HTTPServer, object):
+    def __init__(self, server_address, RequestHandlerClass, key_queue):
+        super(LoginServer, self).__init__(server_address, RequestHandlerClass)
+        self.key_queue = key_queue
+
+
+class LoginHttpRequestHandler(BaseHTTPRequestHandler):
+    def do_OPTIONS(self):
+        self.send_response(200, "ok")
+        self.send_header('Access-Control-Allow-Origin', '*')
+        self.send_header('Access-Control-Allow-Methods', 'GET, OPTIONS')
+        self.end_headers()
+
+    def do_GET(self):
+        params = urlparse.parse_qs(urlparse.urlparse(self.path).query)
+        key = params.get('apikey')
+        if not key:
+            self.send_response(400)
+            return
+
+        self.server.key_queue.put(key[0])
+
+        self.send_response(200)
+        self.send_header('Access-Control-Allow-Origin', '*')
+        self.send_header('Access-Control-Allow-Methods', 'GET, OPTIONS')
+        self.send_header('Content-type', 'text/html')
+        self.end_headers()
+        self.wfile.write("OK".encode('utf-8'))
+
+    def log_message(self, fmt, *args):
+        return
+
+
+@contextlib.contextmanager
+def get_free_port():
+    try:
+        s = socket.socket()
+        s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+        s.bind(('localhost', 0))
+        yield s.getsockname()[1]
+        s.close()
+    except socket.error:
+        yield None
+
+
+def wait_for_apikey():
+    floyd_logger.info('Waiting for login from browser...')
+
+    key_queue = Queue()
+    with get_free_port() as port:
+        if not port:
+            floyd_logger.error("Failed to allocate TCP port for automatic login.")
+            return
+        server = LoginServer(('', port), LoginHttpRequestHandler, key_queue)
+
+    t = threading.Thread(
+        target=server.serve_forever)
+    t.daemon = True
+    t.start()
+
+    cli_host = 'http://127.0.0.1'
+    url = '%s/cli_login?callback=%s:%s' % (floyd.floyd_web_host, cli_host, port)
+    subprocess.check_output(
+        [sys.executable, '-m', 'webbrowser', url], stderr=subprocess.STDOUT)
+
+    wait_timeout_sec = 0.5
+    wait_cnt = 0
+    while True:
+        if wait_cnt > 60:
+            floyd_logger.error("Failed to get login info from browser, please login manually by creating login key at %s/settings/apikey.", floyd.floyd_web_host)
+            server.shutdown()
+            sys.exit(1)
+        try:
+            apikey = key_queue.get(timeout=wait_timeout_sec)
+            break
+        except QueueEmpty:
+            wait_cnt += 1
+
+    server.shutdown()
+    return apikey
+
+
+def has_browser():
+    try:
+        webbrowser.get()
+        return True
+    except webbrowser.Error:
+        return False