silviud · October 22, 2024 17:56 · Oct 22, 2024
diff --git a/databricks-auth-token.py b/databricks-auth-token.py
@@ -0,0 +1,87 @@
+# https://docs.databricks.com/en/dev-tools/auth/oauth-u2m.html#generate-oauth-pair
+# the auth is handled via azure cli
+# you need to issue `az login` and set the sub name|id where the ws is attached.
+
+import cmd
+import uuid, hashlib, base64
+
+
+def generate_codes():
+    # Generate a UUID.
+    uuid1 = uuid.uuid4()
+
+    # Convert the UUID to a string.
+    uuid_str1 = str(uuid1).upper()
+
+    # Create the code verifier.
+    code_verifier = uuid_str1 + "-" + uuid_str1
+
+    # Create the code challenge based on the code verifier.
+    code_challenge = base64.urlsafe_b64encode(hashlib.sha256(code_verifier.encode()).digest()).decode('utf-8')
+
+    # Remove all padding from the code challenge.
+    code_challenge = code_challenge.replace('=', '')
+
+    # Print the code verifier and the code challenge.
+    # Use these in your calls to manually generate
+    # access tokens for OAuth U2M authentication.
+    # print(f"code_verifier:  {code_verifier}")
+    # print(f"code_challenge: {code_challenge}")
+
+    return code_verifier, code_challenge
+
+
+def build_authorization_code_url(workspace_url, code_challenge, redirect_url='http://localhost:8020', state_str='databricks-cli'):
+    return f"https://{workspace_url}/oidc/v1/authorize?client_id=databricks-cli&redirect_uri={redirect_url}&response_type=code&state={state_str}&code_challenge={code_challenge}&code_challenge_method=S256&scope=all-apis+offline_access"
+
+
+def build_access_token_url(workspace_url, authorization_code, code_verifier, redirect_url='http://localhost:8020'):
+    return f'curl --request POST https://{workspace_url}/oidc/v1/token --data "client_id=databricks-cli" --data "grant_type=authorization_code" --data "scope=all-apis offline_access" --data "redirect_uri={redirect_url}" --data "code_verifier={code_verifier}" --data "code={authorization_code}"'
+
+
+class DatabricksCmd(cmd.Cmd):
+
+    intro = 'This program obtains an access token from Databricks. Type help or ? to list commands'
+    prompt = '(db-cli) '
+
+    workspace_url = None
+    redirect_url  = 'http://localhost:8020'
+    state_str     = 'databricks-cli'
+    code_verifier, code_challenge = generate_codes()
+
+    def do_set_host_url(self, arg):
+        """ Set the host url - Can be account or workspace """
+        print(arg)
+        self.workspace_url = arg
+
+    def do_get_settings(self, arg):
+        """ Print current host """
+        print(f'workspace_url={self.workspace_url}')
+        print(f'redirect_url={self.redirect_url}')
+        print(f'state_str={self.state_str}')
+        print(f'code_verifier={self.code_verifier}')
+        print(f'code_challenge={self.code_challenge}')
+
+    def do_get_authorization_code_url(self, arg):
+        """ Constructs an authorization url that
+        you can paste in your browser to obtain the authorization code.
+        """
+        print(build_authorization_code_url(self.workspace_url, self.code_challenge))
+
+    def do_get_authentication_token_url(self, arg):
+        """ Constructs url to obtain access token
+        Paste in the authorization code obtain from browser.
+        The code starts with 'dcod*'.
+        """
+        authorization_code = arg  # TODO validate it starts with dc
+        print(build_access_token_url(self.workspace_url, authorization_code, self.code_verifier))
+
+    def do_bye(self, arg):
+        'Close this program.'
+        print('Bye')
+        return True
+
+
+if __name__ == '__main__':
+
+    DatabricksCmd().cmdloop()
No results found