suite-py 1.41.3__py3-none-any.whl → 1.41.4__py3-none-any.whl

suite_py/commands/login.py

@@ -1,180 +1,11 @@
-import base64
-import hashlib
-import logging
-import secrets
-import sys
-import threading
-import urllib
-import webbrowser
-from time import sleep
-
-import requests
-from flask import Flask, render_template, request
-from werkzeug.serving import make_server
-
-from suite_py.lib import logger
 from suite_py.lib import metrics
-
-# Global vars for comunication between flask thread and main cli function
-received_callback = None
-code = None
-error_message = None
-received_state = None
-
-# Global Flask App
-app = Flask(__name__, template_folder="templates")
-log = logging.getLogger("werkzeug")
-log.disabled = True
-
-
-@app.route("/callback")
-def callback():
-    """# pylint: disable-next=missing-timeout
-    The callback is invoked after a completed login attempt (succesful or otherwise).
-    It sets global variables with the auth code or error messages, then sets the
-    polling flag received_callback.
-    :return:
-    """
-    global received_callback, code, error_message, received_state
-    error_message = None
-    code = None
-    if "error" in request.args:
-        error_message = request.args["error"] + ": " + request.args["error_description"]
-    else:
-        code = request.args["code"]
-    received_state = request.args["state"]
-    received_callback = True
-    return render_template("login.html")
-
-
-class ServerThread(threading.Thread):
-    """
-    The Flask server is done this way to allow shutting down after a single request has been received.
-    """
-
-    def __init__(self, app):
-        threading.Thread.__init__(self)
-        self.srv = make_server("127.0.0.1", 5000, app)
-        self.ctx = app.app_context()
-        self.ctx.push()
-
-    def run(self):
-        # logger.debug('starting server')
-        self.srv.serve_forever()
-
-    def shutdown(self):
-        self.srv.shutdown()
+from suite_py.lib.handler.okta_handler import Okta
 
 
 class Login:
-    def __init__(self, config):
-
-        try:
-            if not config.okta["client_id"] or not config.okta["base_url"]:
-                self.usage()
-                sys.exit(-1)
-        except AttributeError:
-            self.usage()
-            sys.exit(-1)
-
-        self._config = config
-        self.okta_client_id = config.okta["client_id"]
-        self.base_url = config.okta["base_url"]
-        self.okta_scope = "openid"
-        self.redirect_uri = "http://127.0.0.1:5000/callback"
-
-    def usage(self):
-        logger.warning("Unable to login: missing config")
-        logger.warning(
-            "please check docs: https://github.com/primait/suite_py/blob/master/README.md"
-        )
-
-    def url_encode_no_padding(self, byte_data):
-        """
-        Safe encoding handles + and /, and also replace = with nothing
-        :param byte_data:
-        :return:
-        """
-        return base64.urlsafe_b64encode(byte_data).decode("utf-8").replace("=", "")
-
-    def generate_challenge(self, a_verifier):
-        return self.url_encode_no_padding(hashlib.sha256(a_verifier.encode()).digest())
+    def __init__(self, config, tokens):
+        self._okta = Okta(config, tokens)
 
     @metrics.command("login")
     def run(self):
-        global received_callback
-
-        # from https://auth0.com/docs/flows/add-login-using-the-authorization-code-flow-with-pkce
-        # Step1: Create code verifier: Generate a code_verifier that will be sent to Auth0 to request tokens.
-        verifier = self.url_encode_no_padding(secrets.token_bytes(32))
-        # Step2: Create code challenge: Generate a code_challenge from the code_verifier that will be sent to Auth0 to request an authorization_code.
-        challenge = self.generate_challenge(verifier)
-        state = self.url_encode_no_padding(secrets.token_bytes(32))
-        client_id = self.okta_client_id
-        base_url = self.base_url
-        scope = self.okta_scope
-        redirect_uri = self.redirect_uri
-        #
-        # We generate a nonce (state) that is used to protect against attackers invoking the callback
-        url = f"{base_url}/authorize?"
-        url_parameters = {
-            "scope": scope,
-            "response_type": "code",
-            "redirect_uri": redirect_uri,
-            "client_id": client_id,
-            "code_challenge": challenge.replace("=", ""),
-            "code_challenge_method": "S256",
-            "state": state,
-        }
-        url = url + urllib.parse.urlencode(url_parameters)
-
-        # Step3: Authorize user: Request the user's authorization and redirect back to your app with an authorization_code.
-        # Open the browser window to the login url
-        # Start the server
-        # Poll til the callback has been invoked
-        received_callback = False
-        logger.info(
-            "A browser tab should've opened. If not manually navigate to: " + url
-        )
-        webbrowser.open(url)
-        server = ServerThread(app)
-        server.start()
-        while not received_callback:
-            sleep(1)
-        server.shutdown()
-
-        if state != received_state:
-            logger.error(
-                "Error: session replay or similar attack in progress. Please log out of all connections."
-            )
-            sys.exit(1)
-
-        if error_message:
-            logger.error("An error occurred:")
-            logger.error(error_message)
-            sys.exit(1)
-
-        # Step4: Request tokens: Exchange your authorization_code and code_verifier for tokens.
-        url = f"{base_url}/token"
-        headers = {
-            "content-type": "application/x-www-form-urlencoded",
-            "Accept": "application/json",
-        }
-        body = {
-            "grant_type": "authorization_code",
-            "client_id": client_id,
-            "code_verifier": verifier,
-            "code": code,
-            "redirect_uri": redirect_uri,
-        }
-        # pylint: disable-next=missing-timeout
-        r = requests.post(url, headers=headers, data=body)
-        data = r.json()
-        logger.debug(data)
-        if "id_token" in data:
-            token = str(data["id_token"])
-        else:
-            logger.error("id_token not found in okta response")
-            sys.exit(1)
-        logger.info("login succeded")
-        self._config.put_cache(f"{base_url}_token", token)
+        self._okta.login()

suite_py/commands/merge_pr.py

@@ -3,8 +3,7 @@ import sys
 
 from halo import Halo
 
-from suite_py.lib import logger
-from suite_py.lib import metrics
+from suite_py.lib import logger, metrics
 from suite_py.lib.handler import git_handler as git
 from suite_py.lib.handler import prompt_utils
 from suite_py.lib.handler.captainhook_handler import CaptainHook
@@ -16,11 +15,11 @@ from suite_py.lib.symbol import CHECKMARK, CROSSMARK
 
 
 class MergePR:
-    def __init__(self, project, config, tokens):
+    def __init__(self, project, captainhook: CaptainHook, config, tokens):
         self._project = project
         self._config = config
         self._youtrack = YoutrackHandler(config, tokens)
-        self._captainhook = CaptainHook(config, tokens=tokens)
+        self._captainhook = captainhook
         self._git = GitHandler(project, config)
         self._github = GithubHandler(tokens)
         self._drone = DroneHandler(config, tokens, repo=project)

suite_py/commands/open_pr.py

@@ -3,9 +3,7 @@ import sys
 
 from github import GithubException
 
-from suite_py.commands.ask_review import AskReview
-from suite_py.lib import logger
-from suite_py.lib import metrics
+from suite_py.lib import logger, metrics
 from suite_py.lib.handler import prompt_utils
 from suite_py.lib.handler.git_handler import GitHandler, get_commit_logs
 from suite_py.lib.handler.github_handler import GithubHandler
@@ -13,7 +11,8 @@ from suite_py.lib.handler.youtrack_handler import YoutrackHandler
 
 
 class OpenPR:
-    def __init__(self, project, config, tokens):
+    def __init__(self, ask_review, project, config, tokens):
+        self._ask_review = ask_review
         self._project = project
         self._config = config
         self._tokens = tokens
@@ -115,7 +114,7 @@ class OpenPR:
             return
 
         if prompt_utils.ask_confirm("Do you want to insert reviewers?"):
-            AskReview(self._project, self._config, self._tokens).run()
+            self._ask_review.run()
         elif youtrack_id:
             self._detect_and_move_new_pr_with_reviewers(youtrack_id, pr)
 

suite_py/commands/project_lock.py

@@ -3,17 +3,15 @@ import sys
 
 import requests
 
-from suite_py.lib import logger
-from suite_py.lib import metrics
-from suite_py.lib.handler.captainhook_handler import CaptainHook
+from suite_py.lib import logger, metrics
 
 
 class ProjectLock:
-    def __init__(self, project, env, action, config, tokens):
+    def __init__(self, project, env, action, captainhook):
         self._project = project
         self._env = _parse_env(env)
         self._action = action
-        self._captainhook = CaptainHook(config, tokens=tokens)
+        self._captainhook = captainhook
 
     @metrics.command("manage-project-lock")
     def run(self):

suite_py/commands/release.py

@@ -6,11 +6,9 @@ import semver
 from halo import Halo
 
 from suite_py.commands import common
-from suite_py.lib import logger
-from suite_py.lib import metrics
+from suite_py.lib import logger, metrics
 from suite_py.lib.handler import git_handler as git
 from suite_py.lib.handler import prompt_utils
-from suite_py.lib.handler.captainhook_handler import CaptainHook
 from suite_py.lib.handler.changelog_handler import ChangelogHandler
 from suite_py.lib.handler.drone_handler import DroneHandler
 from suite_py.lib.handler.git_handler import GitHandler
@@ -21,7 +19,7 @@ from suite_py.lib.handler.youtrack_handler import YoutrackHandler
 
 class Release:
     # pylint: disable=too-many-instance-attributes
-    def __init__(self, action, project, config, tokens, flags=None):
+    def __init__(self, action, project, captainhook, config, tokens, flags=None):
         self._action = action
         self._project = project
         self._flags = flags
@@ -29,7 +27,7 @@ class Release:
         self._tokens = tokens
         self._changelog_handler = ChangelogHandler()
         self._youtrack = YoutrackHandler(config, tokens)
-        self._captainhook = CaptainHook(config)
+        self._captainhook = captainhook
         self._github = GithubHandler(tokens)
         self._repo = self._github.get_repo(project)
         self._git = GitHandler(project, config)

suite_py/commands/secret.py

@@ -5,8 +5,7 @@ import sys
 
 import yaml
 
-from suite_py.lib import logger
-from suite_py.lib import metrics
+from suite_py.lib import logger, metrics
 from suite_py.lib.handler import prompt_utils
 from suite_py.lib.handler.vault_handler import VaultHandler
 

suite_py/commands/set_token.py

@@ -1,7 +1,6 @@
 import sys
 
-from suite_py.lib import logger
-from suite_py.lib import metrics
+from suite_py.lib import logger, metrics
 from suite_py.lib.handler import prompt_utils
 from suite_py.lib.tokens import Tokens
 

suite_py/commands/status.py

@@ -1,16 +1,15 @@
 # -*- coding: utf-8 -*-
 from halo import Halo
 
-from suite_py.lib import logger
-from suite_py.lib import metrics
+from suite_py.lib import logger, metrics
 from suite_py.lib.handler.captainhook_handler import CaptainHook
 from suite_py.lib.symbol import CHECKMARK, CROSSMARK
 
 
 class Status:
-    def __init__(self, project, config):
+    def __init__(self, project, captainhook: CaptainHook):
         self._project = project
-        self._captainhook = CaptainHook(config)
+        self._captainhook = captainhook
 
     @metrics.command("status")
     def run(self):

suite_py/lib/handler/captainhook_handler.py

@@ -2,6 +2,8 @@
 import requests
 
 from suite_py.lib.handler.github_handler import GithubHandler
+from suite_py.lib.handler.okta_handler import Okta
+from suite_py.__version__ import __version__
 
 
 class UnauthorizedError(Exception):
@@ -10,14 +12,12 @@ class UnauthorizedError(Exception):
 
 
 class CaptainHook:
-    def __init__(self, config, tokens=None):
+    _okta: Okta
+
+    def __init__(self, config, okta: Okta, tokens=None):
         self._baseurl = config.user["captainhook_url"]
         self._timeout = config.user["captainhook_timeout"]
-        self._okta_base_url = config.okta["base_url"]
-        self._okta_token = config.get_cache(f"{self._okta_base_url}_token")
-        self._headers = {
-            "Authorization": f"Bearer {self._okta_token}",
-        }
+        self._okta = okta
 
         if tokens is not None:
             self._github = GithubHandler(tokens)
@@ -64,7 +64,7 @@ class CaptainHook:
     def send_post_request(self, endpoint, data=None, json=None):
         r = requests.post(
             f"{self._baseurl}{endpoint}",
-            headers=self._headers,
+            headers=self._headers(),
             data=data,
             json=json,
             timeout=self._timeout,
@@ -75,7 +75,7 @@ class CaptainHook:
     def send_put_request(self, endpoint, data=None, json=None):
         r = requests.put(
             f"{self._baseurl}{endpoint}",
-            headers=self._headers,
+            headers=self._headers(),
             data=data,
             json=json,
             timeout=self._timeout,
@@ -86,7 +86,7 @@ class CaptainHook:
     def send_get_request(self, endpoint):
         r = requests.get(
             f"{self._baseurl}{endpoint}",
-            headers=self._headers,
+            headers=self._headers(),
             timeout=(2, self._timeout),
         )
 
@@ -103,3 +103,10 @@ class CaptainHook:
 
     def _get_user(self):
         return self._github.get_user().login
+
+    def _headers(self):
+        id_token = self._okta.get_id_token()
+        return {
+            "User-Agent": f"suite-py/{__version__}",
+            "Authorization": f"Bearer {id_token}",
+        }

suite_py/lib/handler/metrics_handler.py

@@ -12,7 +12,8 @@ from suite_py.lib.handler.captainhook_handler import CaptainHook
 
 # Collects metrics and sends them to datadog through captainhook
 class Metrics:
-    def __init__(self, config: Config):
+    def __init__(self, captainhook: CaptainHook, config: Config):
+        self._captainhook = captainhook
         self._config = config
 
     # Emits the command_executed metric
@@ -37,8 +38,11 @@ class Metrics:
             metrics = []
 
         logger.debug(f"creating metric: {metric}")
-        metrics.append(metric)
-        self._config.put_cookie("metrics", metrics)
+        if self._config.user.get("disable_metrics_creation", False):
+            logger.debug("skipping metric creation")
+        else:
+            metrics.append(metric)
+            self._config.put_cookie("metrics", metrics)
 
     # Upload metrics in a detached background process
     def async_upload(self):
@@ -69,8 +73,6 @@ class Metrics:
             self.upload()
 
     def upload(self):
-        captainhook = CaptainHook(self._config)
-
         metrics = self._config.get_cookie("metrics", [])
         # Prevent double uploads
         #
@@ -79,7 +81,7 @@ class Metrics:
         self._config.put_cookie("metrics", [])
         try:
             if len(metrics) != 0:
-                captainhook.send_metrics(metrics)
+                self._captainhook.send_metrics(metrics)
         except Exception:
             # Upload failed, try again later
             self._config.put_cookie("metrics", metrics)

suite_py/lib/handler/okta_handler.py

@@ -0,0 +1,81 @@
+import time
+import typing
+
+from suite_py.lib import logger, oauth
+from suite_py.lib.config import Config
+from suite_py.lib.tokens import Tokens
+
+_SCOPE = "openid offline_access"
+
+
+class Okta:
+    def __init__(self, config: Config, tokens: Tokens) -> None:
+        self._config = config
+        self._tokens = tokens
+
+    def login(self):
+        res = oauth.authorization_code_flow(
+            self._config.okta["client_id"],
+            self._config.okta["base_url"],
+            _SCOPE,
+        )
+
+        self._update_tokens(res)
+
+    def get_id_token(self) -> str:
+        """
+        Returns an id_token, performing a token refresh if needed
+        Raises on an invalid or missing refresh token
+        """
+        return self._get_id_token() or self._refresh()
+
+    def _refresh(self) -> str:
+        logger.debug("Refreshing id_token")
+
+        refresh_token = self._get_refresh_token()
+        if not isinstance(refresh_token, str):
+            raise Exception(
+                "Invalid okta refresh token. Try logging in with `suite_py login`"
+            )
+        res = oauth.do_refresh_token(
+            self._config.okta["client_id"],
+            self._config.okta["base_url"],
+            _SCOPE,
+            refresh_token,
+        )
+
+        return self._update_tokens(res)
+
+    def _update_tokens(self, tokens: oauth.OAuthTokenResponse) -> str:
+        if not tokens.id_token:
+            raise Exception("Okta didn't return a new id_token. This shouldn't happen.")
+        if not tokens.refresh_token:
+            raise Exception(
+                "Okta didn't return a new refresh_token. This shouldn't happen."
+            )
+
+        self._set_refresh_token(tokens.refresh_token)
+
+        expires_at = time.time() + tokens.expires_in
+        self._set_id_token(tokens.id_token, expires_at)
+
+        return tokens.id_token
+
+    def _get_refresh_token(self) -> typing.Optional[str]:
+        return self._tokens.okta().get("refresh_token", None)
+
+    def _set_refresh_token(self, token: str):
+        okta = self._tokens.okta()
+        okta["refresh_token"] = token
+        self._tokens.edit("okta", okta)
+
+    def _get_id_token(self) -> typing.Optional[str]:
+        (token, expiration) = self._tokens.okta().get("id_token", (None, None))
+        if token is not None and time.time() < expiration:
+            return token
+        return None
+
+    def _set_id_token(self, token: str, expiration: float):
+        okta = self._tokens.okta()
+        okta["id_token"] = (token, expiration)
+        self._tokens.edit("okta", okta)

suite_py/lib/logger.py

@@ -1,5 +1,6 @@
 # -*- encoding: utf-8 -*-
 import logging
+
 import logzero
 from logzero import logger as _logger
 

suite_py/lib/metrics.py

@@ -1,6 +1,8 @@
 # -*- encoding: utf-8 -*-
 import functools
 from suite_py.lib.config import Config
+
+from suite_py.lib.handler.captainhook_handler import CaptainHook
 from suite_py.lib.handler.metrics_handler import Metrics
 
 _metrics_handler = None
@@ -15,9 +17,9 @@ def _metrics() -> Metrics:
     )
 
 
-def setup(config: Config):
+def setup(config: Config, captainhook: CaptainHook):
     global _metrics_handler
-    _metrics_handler = Metrics(config)
+    _metrics_handler = Metrics(config=config, captainhook=captainhook)
 
 
 def command_executed(command):