import prefect
from prefect import Client
from prefect import task, Flow


@task(log_stdout=True)
def get_user_info():
    client = Client()
    flow_run_id = prefect.context.get("flow_run_id")
    query = (
        """
            query {
              flow_run(where: { id: {_eq: \""""
        + flow_run_id
        + """\"} })  {
                id
                name
                created_by {
                  username
                  id
                  email
                  first_name
                  last_name
                }
              }
            }
        """
    )
    results = client.graphql(query)
    print(results)


with Flow("user") as flow:
    get_user_info()

import ldap3
import prefect 

class SecurityPrincipal:
    """
    A group of attributes and methods required for authentication and identification
    
    Args:
        - service (str): the name of a server, website or service we wish to authenticate 
            against. It may be a functional URI or merely a label for user identification
        - user (str, optional): the username of a security principal to be authenticated
        - distinguished_name (str, optional): string required to authenticate when using ldap
        - name (str, optional): a label that identifies the principal
        - mail (str, optional): a user or person's email address
    """
    def __init__(
        self,
        service: str,
        user: str = None,
        distinguished_name: str = None,
        name: str = None,
        mail: str = None,
    ):
        self.service = service
        self.user = user
        self.distinguished_name = distinguished_name
        self.name = name
        self.mail = mail
        
    def set_local_secret(self, secret):
        """
        Accepts a string in clear text to use as a secret, maybe a password
        """
        self.secret = secret
        return self
        
    def get_secret(self):
        """
        Returns your clear string secret
        
        This method is expected to be replaced in sub-classes, such as retrieval 
        of a secret from the Windows Credential Manager, or an Azure Key Vault. 
        In such cases the secret attribute set by set_local_secret will take precedence.
        """
        try:
            return self.secret
        except AttributeError:
            raise AttributeError(f'Secret value not set for service {self.service}')
            
    def ldap_search(
        self,
        ldap_user: "SecurityPrincipal",
        search_base: str,
        search_attribute: str = 'mail',
        objectclass: str = 'user',
        returned_attributes: list[str] = ['name'],
    ):
        """
        Performs an ldap search and stores the requested attributes from the first entry found
        
        Args:
            - ldap_user (SecurityPrincipal): for authentication against the DIT
            - search_base (str): the location in the DIT where the search will start
            - search_attribute (str, optional): name of a known attribute value to compare 
                against in the DIT. Default is 'mail'
            - objectclass: str, optional): the class of the object to find. Default is 'user'
            - returned_attributes (list[str], optional): a list of attribute names that will 
                be retrieved and stored from the first entry returned. Default is ['name']
        """
        search_filter = 
            f'(&({search_attribute}={getattr(self, search_attribute)})(objectclass={objectclass}))'
        server = ldap3.Server(
            ldap_user.service, use_ssl=True, get_info=ldap3.ALL
        )
        connection = ldap3.Connection(
            server, ldap_user.distinguished_name, ldap_user.get_secret(), auto_bind=True
        )
        connection.search(
            search_base=search_base,
            search_filter=search_filter,
            attributes=returned_attributes
        )
        first_entry = connection.entries[0]
        for attr in returned_attributes:
            setattr(self, attr, getattr(first_entry, attr))
        return self
        
    def prefect_lookup(
        self,
        source_attribute: str = 'email',
        destination_attribute: str = 'mail',
    ):
        """
        Stores information about a flow run creator from the prefect cloud ui to an attribute
        
        Args:
            - source_attribute (str): name of the key value to be retrieved from prefect cloud 
                ui. Default is 'email'
            - destination_attribute (str): name of the attribute in this class instance to attach 
                the retrieved value to. Default is 'mail'
        """
        client = prefect.Client()
        flow_run_id = prefect.context.get('flow_run_id')
        query = 'query {flow_run(where: { id: {_eq: "' + flow_run_id + '"} }) {created_by {' + source_attribute + '}}}'
        result = client.graphql(query)
        setattr(self, destination_attribute, result['data']['flow_run'][0]['created_by'][source_attribute])
        return self