Move handler back and method we're expecting to override

awscli/botocore/utils.py

@@ -3085,6 +3085,15 @@ def __init__(
             cache = {}
         self._cache = cache
 
+    def fetch_token(
+        self,
+        start_url,
+        force_refresh,
+        registration_scopes,
+        session_name,
+    ):
+        raise NotImplementedError('Must implement fetch_token()')
+
     def _utc_now(self):
         return datetime.datetime.now(tzutc())
 
@@ -3177,15 +3186,6 @@ def __init__(
             sleep = time.sleep
         self._sleep = sleep
 
-    def fetch_token(
-        self,
-        start_url,
-        force_refresh,
-        registration_scopes,
-        session_name,
-    ):
-        raise NotImplementedError('Must implement fetch_token()')
-
     def _register_client(self, session_name, scopes):
         register_kwargs = {
             'clientName': self._generate_client_name(session_name),

awscli/customizations/sso/utils.py

@@ -199,7 +199,7 @@ def __init__(self):
         # We do this so that the request handler can have a reference to this
         # AuthCodeFetcher so that it can pass back the state and auth code
         try:
-            handler = partial(self.OAuthCallbackHandler, self)
+            handler = partial(OAuthCallbackHandler, self)
             self.http_server = HTTPServer(('', 0), handler)
         except socket.error as e:
             raise AuthCodeFetcherError(error_msg=e)
@@ -225,41 +225,42 @@ def set_auth_code_and_state(self, auth_code, state):
         self._state = state
         self._is_done = True
 
-    class OAuthCallbackHandler(BaseHTTPRequestHandler):
-        """HTTP handler to handle OAuth callback requests, extracting
-        the auth code and state parameters, and displaying a page directing
-        the user to return to the CLI.
-        """
-        def __init__(self, auth_code_fetcher, *args, **kwargs):
-            self._auth_code_fetcher = auth_code_fetcher
-            super().__init__(*args, **kwargs)
-
-        def log_message(self, format, *args):
-            # Suppress built-in logging, otherwise it prints
-            # each request to console
-            pass
-
-        def do_GET(self):
-            self.send_response(200)
-            self.end_headers()
-            with open(
-                os.path.join(os.path.dirname(__file__), 'index.html'),
-                'rb',
-            ) as file:
-                self.wfile.write(file.read())
-
-            query_params = parse_qs(urlparse(self.path).query)
-
-            if 'error' in query_params:
-                self._auth_code_fetcher.set_auth_code_and_state(
-                    None,
-                    None,
-                )
-            elif 'code' in query_params and 'state' in query_params:
-                self._auth_code_fetcher.set_auth_code_and_state(
-                    query_params['code'][0],
-                    query_params['state'][0],
-                )
+
+class OAuthCallbackHandler(BaseHTTPRequestHandler):
+    """HTTP handler to handle OAuth callback requests, extracting
+    the auth code and state parameters, and displaying a page directing
+    the user to return to the CLI.
+    """
+    def __init__(self, auth_code_fetcher, *args, **kwargs):
+        self._auth_code_fetcher = auth_code_fetcher
+        super().__init__(*args, **kwargs)
+
+    def log_message(self, format, *args):
+        # Suppress built-in logging, otherwise it prints
+        # each request to console
+        pass
+
+    def do_GET(self):
+        self.send_response(200)
+        self.end_headers()
+        with open(
+            os.path.join(os.path.dirname(__file__), 'index.html'),
+            'rb',
+        ) as file:
+            self.wfile.write(file.read())
+
+        query_params = parse_qs(urlparse(self.path).query)
+
+        if 'error' in query_params:
+            self._auth_code_fetcher.set_auth_code_and_state(
+                None,
+                None,
+            )
+        elif 'code' in query_params and 'state' in query_params:
+            self._auth_code_fetcher.set_auth_code_and_state(
+                query_params['code'][0],
+                query_params['state'][0],
+            )
 
 
 class InvalidSSOConfigError(ConfigurationError):