Headless login

CHANGELOG.D/793.feature

@@ -0,0 +1 @@
+Implement `neuro config login-headless` command.

README.md

@@ -26,6 +26,7 @@
 	* [neuro config](#neuro-config)
 		* [neuro config login](#neuro-config-login)
 		* [neuro config login-with-token](#neuro-config-login-with-token)
+		* [neuro config login-headless](#neuro-config-login-headless)
 		* [neuro config show](#neuro-config-show)
 		* [neuro config show-token](#neuro-config-show-token)
 		* [neuro config docker](#neuro-config-docker)
@@ -725,6 +726,7 @@ Name | Description|
 |---|---|
 | _[neuro config login](#neuro-config-login)_| Log into Neuromation Platform |
 | _[neuro config login\-with-token](#neuro-config-login-with-token)_| Log into Neuromation Platform with token |
+| _[neuro config login-headless](#neuro-config-login-headless)_| Log into Neuromation Platform from non-GUI server environment |
 | _[neuro config show](#neuro-config-show)_| Print current settings |
 | _[neuro config show-token](#neuro-config-show-token)_| Print current authorization token |
 | _[neuro config docker](#neuro-config-docker)_| Configure docker client for working with platform registry |
@@ -771,6 +773,25 @@ Name | Description|
 
 
 
+### neuro config login-headless
+
+Log into Neuromation Platform from non-GUI server environment.<br/><br/>URL is a platform entrypoint URL.<br/><br/>The command works similar to "neuro login" but instead of opening a browser<br/>for performing OAuth registration prints an URL that should be open on guest<br/>host.<br/><br/>Then user inputs a code displayed in a browser after successful login back<br/>in neuro command to finish the login process.
+
+**Usage:**
+
+```bash
+neuro config login-headless [OPTIONS] [URL]
+```
+
+**Options:**
+
+Name | Description|
+|----|------------|
+|_--help_|Show this message and exit.|
+
+
+
+
 ### neuro config show
 
 Print current settings.

neuromation/api/__init__.py

@@ -1,7 +1,5 @@
-import sys
 from pathlib import Path
-from types import TracebackType
-from typing import Any, Awaitable, Coroutine, Generator, Optional, Type
+from typing import Awaitable, Callable, Optional
 
 import aiohttp
 from yarl import URL
@@ -40,26 +38,7 @@
 from .parsing_utils import ImageNameParser
 from .storage import FileStatus, FileStatusType
 from .users import Action, Permission, SharedPermission
-
-
-if sys.version_info >= (3, 7):
-    from typing import AsyncContextManager
-else:
-    from typing import Generic, TypeVar
-
-    _T = TypeVar("_T")
-
-    class AsyncContextManager(Generic[_T]):
-        async def __aenter__(self) -> _T:
-            pass  # pragma: no cover
-
-        async def __aexit__(
-            self,
-            exc_type: Optional[Type[BaseException]],
-            exc: Optional[BaseException],
-            tb: Optional[TracebackType],
-        ) -> Optional[bool]:
-            pass  # pragma: no cover
+from .utils import _ContextManager
 
 
 __all__ = (
@@ -103,50 +82,24 @@ async def __aexit__(
 )
 
 
-class _ContextManager(Awaitable[Client], AsyncContextManager[Client]):
-
-    __slots__ = ("_coro", "_client")
-
-    def __init__(self, coro: Coroutine[Any, Any, Client]) -> None:
-        self._coro = coro
-        self._client: Optional[Client] = None
-
-    def __await__(self) -> Generator[Any, None, Client]:
-        return self._coro.__await__()
-
-    async def __aenter__(self) -> Client:
-        self._client = await self._coro
-        assert self._client is not None
-        return self._client
-
-    async def __aexit__(
-        self,
-        exc_type: Optional[Type[BaseException]],
-        exc: Optional[BaseException],
-        tb: Optional[TracebackType],
-    ) -> Optional[bool]:
-        assert self._client is not None
-        await self._client.close()
-        return None
-
-
 def get(
     *, path: Optional[Path] = None, timeout: aiohttp.ClientTimeout = DEFAULT_TIMEOUT
-) -> _ContextManager:
-    return _ContextManager(_get(path, timeout))
+) -> _ContextManager[Client]:
+    return _ContextManager[Client](_get(path, timeout))
 
 
 async def _get(path: Optional[Path], timeout: aiohttp.ClientTimeout) -> Client:
     return await Factory(path).get(timeout=timeout)
 
 
 async def login(
+    show_browser_cb: Callable[[URL], Awaitable[None]],
     *,
     url: URL = DEFAULT_API_URL,
     path: Optional[Path] = None,
     timeout: aiohttp.ClientTimeout = DEFAULT_TIMEOUT
 ) -> None:
-    await Factory(path).login(url=url, timeout=timeout)
+    await Factory(path).login(show_browser_cb, url=url, timeout=timeout)
 
 
 async def login_with_token(
@@ -159,5 +112,15 @@ async def login_with_token(
     await Factory(path).login_with_token(token, url=url, timeout=timeout)
 
 
+async def login_headless(
+    get_auth_code_cb: Callable[[URL], Awaitable[str]],
+    *,
+    url: URL = DEFAULT_API_URL,
+    path: Optional[Path] = None,
+    timeout: aiohttp.ClientTimeout = DEFAULT_TIMEOUT
+) -> None:
+    await Factory(path).login_headless(get_auth_code_cb, url=url, timeout=timeout)
+
+
 async def logout(*, path: Optional[Path] = None) -> None:
     await Factory(path).logout()

neuromation/api/client.py

@@ -1,9 +1,7 @@
-import ssl
 from types import TracebackType
 from typing import Optional, Type
 
 import aiohttp
-import certifi
 
 from .config import _Config
 from .core import DEFAULT_TIMEOUT, _Core
@@ -16,15 +14,17 @@
 
 class Client(metaclass=NoPublicConstructor):
     def __init__(
-        self, config: _Config, *, timeout: aiohttp.ClientTimeout = DEFAULT_TIMEOUT
+        self,
+        connector: aiohttp.BaseConnector,
+        config: _Config,
+        *,
+        timeout: aiohttp.ClientTimeout = DEFAULT_TIMEOUT
     ) -> None:
         config.check_initialized()
         self._config = config
-        self._ssl_context = ssl.SSLContext()
-        self._ssl_context.load_verify_locations(capath=certifi.where())
-        self._connector = aiohttp.TCPConnector(ssl=self._ssl_context)
+        self._connector = connector
         self._core = _Core(
-            self._connector, self._config.url, self._config.auth_token.token, timeout
+            connector, self._config.url, self._config.auth_token.token, timeout
         )
         self._jobs = Jobs._create(self._core, self._config)
         self._storage = Storage._create(self._core, self._config)

neuromation/api/config_factory.py

@@ -1,10 +1,13 @@
+import asyncio
 import os
+import ssl
 import sys
 from dataclasses import replace
 from pathlib import Path
-from typing import Any, Dict, Optional
+from typing import Any, Awaitable, Callable, Dict, Optional
 
 import aiohttp
+import certifi
 import yaml
 from yarl import URL
 
@@ -13,12 +16,15 @@
 from .core import DEFAULT_TIMEOUT
 from .login import (
     AuthNegotiator,
+    HeadlessNegotiator,
     RunPreset,
     _AuthConfig,
     _AuthToken,
     _ClusterConfig,
     get_server_config,
+    refresh_token,
 )
+from .utils import _ContextManager
 
 
 WIN32 = sys.platform == "win32"
@@ -27,6 +33,17 @@
 DEFAULT_API_URL = URL("https://staging.neu.ro/api/v1")
 
 
+def _make_connector() -> _ContextManager[aiohttp.TCPConnector]:
+    return _ContextManager[aiohttp.TCPConnector](__make_connector())
+
+
+async def __make_connector() -> aiohttp.TCPConnector:
+    ssl_context = ssl.SSLContext()
+    ssl_context.load_verify_locations(capath=certifi.where())
+    connector = aiohttp.TCPConnector(ssl=ssl_context)
+    return connector
+
+
 class ConfigError(RuntimeError):
     pass
 
@@ -39,26 +56,70 @@ def __init__(self, path: Optional[Path] = None) -> None:
 
     async def get(self, *, timeout: aiohttp.ClientTimeout = DEFAULT_TIMEOUT) -> Client:
         config = self._read()
-        new_token = await self._refresh_auth_token(config)
+        connector = await _make_connector()
+        try:
+            new_token = await refresh_token(
+                connector, config.auth_config, config.auth_token, timeout
+            )
+        except asyncio.CancelledError:
+            await connector.close()
+            raise
+        except Exception:
+            await connector.close()
+            raise
         if new_token != config.auth_token:
             new_config = replace(config, auth_token=new_token)
             self._save(new_config)
-            return Client._create(new_config, timeout=timeout)
-        return Client._create(config, timeout=timeout)
+            return Client._create(connector, new_config, timeout=timeout)
+        return Client._create(connector, config, timeout=timeout)
 
     async def login(
         self,
+        show_browser_cb: Callable[[URL], Awaitable[None]],
+        *,
+        url: URL = DEFAULT_API_URL,
+        timeout: aiohttp.ClientTimeout = DEFAULT_TIMEOUT,
+    ) -> None:
+        if self._path.exists():
+            raise ConfigError(f"Config file {self._path} already exists. Please logout")
+        async with _make_connector() as connector:
+            config_unauthorized = await get_server_config(connector, url)
+            negotiator = AuthNegotiator(
+                connector, config_unauthorized.auth_config, show_browser_cb, timeout
+            )
+            auth_token = await negotiator.refresh_token()
+
+            config_authorized = await get_server_config(
+                connector, url, token=auth_token.token
+            )
+        config = _Config(
+            auth_config=config_authorized.auth_config,
+            auth_token=auth_token,
+            cluster_config=config_authorized.cluster_config,
+            pypi=_PyPIVersion.create_uninitialized(),
+            url=url,
+        )
+        self._save(config)
+
+    async def login_headless(
+        self,
+        get_auth_code_cb: Callable[[URL], Awaitable[str]],
         *,
         url: URL = DEFAULT_API_URL,
         timeout: aiohttp.ClientTimeout = DEFAULT_TIMEOUT,
     ) -> None:
         if self._path.exists():
             raise ConfigError(f"Config file {self._path} already exists. Please logout")
-        config_unauthorized = await get_server_config(url)
-        negotiator = AuthNegotiator(config_unauthorized.auth_config)
-        auth_token = await negotiator.refresh_token()
+        async with _make_connector() as connector:
+            config_unauthorized = await get_server_config(connector, url)
+            negotiator = HeadlessNegotiator(
+                connector, config_unauthorized.auth_config, get_auth_code_cb, timeout
+            )
+            auth_token = await negotiator.refresh_token()
 
-        config_authorized = await get_server_config(url, token=auth_token.token)
+            config_authorized = await get_server_config(
+                connector, url, token=auth_token.token
+            )
         config = _Config(
             auth_config=config_authorized.auth_config,
             auth_token=auth_token,
@@ -77,7 +138,8 @@ async def login_with_token(
     ) -> None:
         if self._path.exists():
             raise ConfigError(f"Config file {self._path} already exists. Please logout")
-        server_config = await get_server_config(url, token=token)
+        async with _make_connector() as connector:
+            server_config = await get_server_config(connector, url, token=token)
         config = _Config(
             auth_config=server_config.auth_config,
             auth_token=_AuthToken.create_non_expiring(token),
@@ -134,6 +196,7 @@ def _serialize_auth_config(self, auth_config: _AuthConfig) -> Dict[str, Any]:
             "token_url": str(auth_config.token_url),
             "client_id": auth_config.client_id,
             "audience": auth_config.audience,
+            "headless_callback_url": str(auth_config.headless_callback_url),
             "success_redirect_url": success_redirect_url,
             "callback_urls": [str(u) for u in auth_config.callback_urls],
         }
@@ -171,6 +234,7 @@ def _deserialize_auth_config(self, payload: Dict[str, Any]) -> _AuthConfig:
             token_url=URL(auth_config["token_url"]),
             client_id=auth_config["client_id"],
             audience=auth_config["audience"],
+            headless_callback_url=auth_config["headless_callback_url"],
             success_redirect_url=success_redirect_url,
             callback_urls=tuple(URL(u) for u in auth_config.get("callback_urls", [])),
         )
@@ -204,10 +268,6 @@ def _deserialize_auth_token(self, payload: Dict[str, Any]) -> _AuthToken:
             refresh_token=auth_payload["refresh_token"],
         )
 
-    async def _refresh_auth_token(self, config: _Config) -> _AuthToken:
-        auth_negotiator = AuthNegotiator(config=config.auth_config)
-        return await auth_negotiator.refresh_token(config.auth_token)
-
     def _save(self, config: _Config) -> None:
         payload: Dict[str, Any] = dict()
         payload["url"] = str(config.url)

neuromation/api/core.py

@@ -45,7 +45,7 @@ class _Core:
 
     def __init__(
         self,
-        connector: aiohttp.TCPConnector,
+        connector: aiohttp.BaseConnector,
         base_url: URL,
         token: str,
         timeout: aiohttp.ClientTimeout,
@@ -69,7 +69,7 @@ def __init__(
         }
 
     @property
-    def connector(self) -> aiohttp.TCPConnector:
+    def connector(self) -> aiohttp.BaseConnector:
         return self._connector
 
     @property