Fix pre-commit issues

thinkst · Jan 22, 2024 · 37c763b · 37c763b
1 parent 7e4f463
commit 37c763b
Show file tree

Hide file tree

Showing 3 changed files with 82 additions and 45 deletions.
diff --git a/aws-css-token-infra/CSSClonedSiteCFFunc/index.js b/aws-css-token-infra/CSSClonedSiteCFFunc/index.js
@@ -10,7 +10,7 @@ function handler(event) {
     var referer = '';
     if ('referer' in event.request.headers)
         referer = event.request.headers.referer.value;
-    
+
     if (referer == '' || referer.indexOf(expected_referrer) >= 0) { // Happy case where the referer matches
         var response = {
             statusCode: 200,
@@ -33,4 +33,4 @@ function handler(event) {
         }
     };
     return response;
-}
+}
diff --git a/canarytokens/azure_css.py b/canarytokens/azure_css.py
@@ -3,68 +3,83 @@
 from requests import Response, get, put, delete, post
 from time import sleep
 from cssutils.css import CSSStyleRule
-import logging, cssutils
+import logging
+import cssutils
 
 frontend_settings = FrontendSettings()
 
 BearerToken = str
 
+
 def _auth_to_tenant(tenant_id: str) -> BearerToken:
     """
     Tenant that the client app has permissions to, returns a Graph API Bearer token
     """
-    cred = ClientSecretCredential(tenant_id, frontend_settings.AZUREAPP_ID, frontend_settings.AZUREAPP_SECRET)
-    return cred.get_token('.default').token
+    cred = ClientSecretCredential(
+        tenant_id, frontend_settings.AZUREAPP_ID, frontend_settings.AZUREAPP_SECRET
+    )
+    return cred.get_token(".default").token
+
 
 def _check_if_custom_branding(token: BearerToken, tenant_id: str) -> bool:
     """
     Checks to see if a tenant has custom branding (and if it's not safe to install our custom CSS)
     Reference: https://learn.microsoft.com/en-us/graph/api/organizationalbranding-get
     Returns: True if there is branding present, false otherwise
     """
-    headers = {
-        'Accept-Language': '0',
-        'Authorization': 'Bearer ' + token
-    }
-    res: Response = get(f"https://graph.microsoft.com/v1.0/organization/{tenant_id}/branding", headers=headers)
+    headers = {"Accept-Language": "0", "Authorization": "Bearer " + token}
+    res: Response = get(
+        f"https://graph.microsoft.com/v1.0/organization/{tenant_id}/branding",
+        headers=headers,
+    )
     # This API returns 404 if there is no corporate branding configured
     return res.status_code != 404
 
-def _check_existing_body_background(css : str) -> bool:
-    '''
+
+def _check_existing_body_background(css: str) -> bool:
+    """
     Parses an existing CSS file to see if there is a conflicting rule
     Returns True if there is a conflicting rule, False otherwise
-    '''
+    """
     css_rules: list[CSSStyleRule] = cssutils.parseString(css)
     for rule in css_rules:
-        if rule.selectorText == 'body' and 'background' in rule.style.cssText:
+        if rule.selectorText == "body" and "background" in rule.style.cssText:
             return True
     return False
 
-def _check_if_can_install_custom_css(token: BearerToken, tenant_id: str) -> tuple[bool, str]:
+
+def _check_if_can_install_custom_css(
+    token: BearerToken, tenant_id: str
+) -> tuple[bool, str]:
     """
     Checks to see if a tenant has custom branding (and if it's not safe to install our custom CSS)
     References:
         - https://learn.microsoft.com/en-us/graph/api/organizationalbranding-get
         - https://learn.microsoft.com/en-us/graph/api/resources/organizationalbrandingproperties
     Returns: A tuple of (True, css) if we can safely install a custom CSS, False otherwise
     """
-    headers = {
-        'Accept-Language': '0',
-        'Authorization': 'Bearer ' + token
-    }
-    res: Response = get(f"https://graph.microsoft.com/v1.0/organization/{tenant_id}/branding", headers=headers)
-    if res.status_code == 404: # There is no branding at all
-        return (True, '')
-    if res.status_code != 200: # Other error
-        return (False, '')
-    if res.json().get('customCSSRelativeUrl') is None: # Is there another CSS? If not then we can install!
-        return (True, '')
+    headers = {"Accept-Language": "0", "Authorization": "Bearer " + token}
+    res: Response = get(
+        f"https://graph.microsoft.com/v1.0/organization/{tenant_id}/branding",
+        headers=headers,
+    )
+    if res.status_code == 404:  # There is no branding at all
+        return (True, "")
+    if res.status_code != 200:  # Other error
+        return (False, "")
+    if (
+        res.json().get("customCSSRelativeUrl") is None
+    ):  # Is there another CSS? If not then we can install!
+        return (True, "")
     # There is an existing CSS, let's check for compatiblity
-    res: Response = get(f"https://graph.microsoft.com/v1.0/organization/{tenant_id}/branding/localizations/0/customCSS", headers=headers)
+    res: Response = get(
+        f"https://graph.microsoft.com/v1.0/organization/{tenant_id}/branding/localizations/0/customCSS",
+        headers=headers,
+    )
     if res.status_code == 200:
         return (not _check_existing_body_background(res.text), res.text)
-    return (False, '')
+    return (False, "")
+
 
 def _install_custom_css(token: BearerToken, tenant_id: str, css: str) -> bool:
     """
@@ -74,33 +89,45 @@ def _install_custom_css(token: BearerToken, tenant_id: str, css: str) -> bool:
         - https://learn.microsoft.com/en-us/graph/api/organizationalbranding-post-localizations
     Returns: True if successful, False otherwise
     """
-    headers = {
-        'Authorization': 'Bearer ' + token,
-        'Accept-Language': '0'
-    }
-    if not _check_if_custom_branding(token, tenant_id): 
+    headers = {"Authorization": "Bearer " + token, "Accept-Language": "0"}
+    if not _check_if_custom_branding(token, tenant_id):
         # If we need to create a default OrganizationalBranding object, we set a blank string to a single space to create it
         logging.error("Creating a new default organizationalBranding object...")
-        res: Response = post(f"https://graph.microsoft.com/v1.0/organization/{tenant_id}/branding/localizations/", json={"usernameHintText": " "}, headers=headers)
+        res: Response = post(
+            f"https://graph.microsoft.com/v1.0/organization/{tenant_id}/branding/localizations/",
+            json={"usernameHintText": " "},
+            headers=headers,
+        )
         if res.status_code != 201:
-            logging.error(f"Unable to create OrganizationalBranding object: {res.status_code} - {res.text}")
+            logging.error(
+                f"Unable to create OrganizationalBranding object: {res.status_code} - {res.text}"
+            )
             return False
-        sleep(5) # Give the Graph API a second to recognize it's built
-
-    headers['Content-Type'] = 'text/css'
-    res: Response = put(f"https://graph.microsoft.com/v1.0/organization/{tenant_id}/branding/localizations/0/customCSS", data=css.encode(), headers=headers)
+        sleep(5)  # Give the Graph API a second to recognize it's built
+
+    headers["Content-Type"] = "text/css"
+    res: Response = put(
+        f"https://graph.microsoft.com/v1.0/organization/{tenant_id}/branding/localizations/0/customCSS",
+        data=css.encode(),
+        headers=headers,
+    )
     if res.status_code != 204:
         logging.error(f"Unable to add customCSS: {res.status_code} - {res.text}")
     return res.status_code == 204
 
+
 def _delete_self(token: BearerToken) -> bool:
     """
     Tries to delete itself from the tenant to reduce risk
     Returns: True is successful, False otherwise
     """
-    res: Response = delete(f"https://graph.microsoft.com/v1.0/servicePrincipals(appId='{frontend_settings.AZUREAPP_ID}')", headers={'Authorization': 'Bearer ' + token})
+    res: Response = delete(
+        f"https://graph.microsoft.com/v1.0/servicePrincipals(appId='{frontend_settings.AZUREAPP_ID}')",
+        headers={"Authorization": "Bearer " + token},
+    )
     return res.status_code == 204
 
+
 def install_azure_css(tenant_id: str, css: str) -> tuple[bool, str]:
     """
     Main business logic function to install the Azure CSS token into the tenant
@@ -110,12 +137,22 @@ def install_azure_css(tenant_id: str, css: str) -> tuple[bool, str]:
     token = _auth_to_tenant(tenant_id)
     (check, existing_css) = _check_if_can_install_custom_css(token, tenant_id)
     if not check:
-        return (False, f"Installation failed: your tenant already has a conflicting custom CSS, please manually add the CSS to your portal branding.")
+        return (
+            False,
+            "Installation failed: your tenant already has a conflicting custom CSS, please manually add the CSS to your portal branding.",
+        )
     if not _install_custom_css(token, tenant_id, existing_css + css):
         # Might as well remove ourselves anyways
         _delete_self(token)
-        return (False, f"Installation failed: Unable to automatically install the CSS, please manually add the CSS to your portal branding.")
+        return (
+            False,
+            "Installation failed: Unable to automatically install the CSS, please manually add the CSS to your portal branding.",
+        )
     _delete_self(token)
-    return (True, "Successfully installed the CSS into your Azure tenant. Please wait for a few minutes for the changes to propogate; no further action is needed.")
+    return (
+        True,
+        "Successfully installed the CSS into your Azure tenant. Please wait for a few minutes for the changes to propogate; no further action is needed.",
+    )
+
 
 # Other reference that's of note but not linked to from the other Graph API docs: https://learn.microsoft.com/en-us/graph/api/organizationalbrandinglocalization-delete
diff --git a/canarytokens/canarydrop.py b/canarytokens/canarydrop.py
@@ -326,8 +326,8 @@ def get_cloned_site_javascript(self, force_https: bool = False):
             """
         )
         return clonedsite_js
-    
-    def get_cloned_site_css(self, cf_url : str):
+
+    def get_cloned_site_css(self, cf_url: str):
         token_val = self.canarytoken.value()
         expected_referrer = quote(b64encode(self.expected_referrer.encode()).decode())
         clonedsite_css = textwrap.dedent(