aimodelshare 0.3.7__py3-none-any.whl

aimodelshare/auth.py

@@ -0,0 +1,163 @@
+"""
+Authentication and identity management helpers for aimodelshare.
+
+Provides unified authentication around Cognito IdToken (JWT_AUTHORIZATION_TOKEN),
+with backward compatibility for legacy AWS_TOKEN.
+"""
+
+import os
+import warnings
+import logging
+from typing import Optional, Dict, Any
+import json
+
+logger = logging.getLogger("aimodelshare.auth")
+
+try:
+    import jwt
+except ImportError:
+    jwt = None
+    logger.warning("PyJWT not installed. JWT decode functionality will be limited.")
+
+
+def get_primary_token() -> Optional[str]:
+    """
+    Get the primary authentication token from environment variables.
+    
+    Prefers JWT_AUTHORIZATION_TOKEN over legacy AWS_TOKEN.
+    Issues a deprecation warning if only AWS_TOKEN is present.
+    
+    Returns:
+        Optional[str]: The authentication token, or None if not found
+    """
+    jwt_token = os.getenv('JWT_AUTHORIZATION_TOKEN')
+    if jwt_token:
+        return jwt_token
+    
+    aws_token = os.getenv('AWS_TOKEN')
+    if aws_token:
+        warnings.warn(
+            "Using legacy AWS_TOKEN environment variable. "
+            "Please migrate to JWT_AUTHORIZATION_TOKEN. "
+            "AWS_TOKEN support will be deprecated in a future release.",
+            DeprecationWarning,
+            stacklevel=2
+        )
+        return aws_token
+    
+    return None
+
+
+def get_identity_claims(token: Optional[str] = None, verify: bool = False) -> Dict[str, Any]:
+    """
+    Extract identity claims from a JWT token.
+    
+    Args:
+        token: JWT token string. If None, uses get_primary_token()
+        verify: If True, performs signature verification (requires JWKS endpoint)
+                Currently defaults to False as JWKS verification is future work
+    
+    Returns:
+        Dict containing identity claims:
+        - sub: Subject (user ID)
+        - email: User email
+        - cognito:username: Username (if present)
+        - iss: Issuer
+        - principal: Derived principal identifier
+    
+    Raises:
+        ValueError: If token is invalid or missing
+        RuntimeError: If PyJWT is not installed
+    
+    Note:
+        This currently performs unverified decode as JWKS signature verification
+        is planned for future work. Do not use in production security-critical
+        contexts without implementing signature verification.
+    """
+    if token is None:
+        token = get_primary_token()
+    
+    if not token:
+        raise ValueError("No authentication token available")
+    
+    if jwt is None:
+        raise RuntimeError("PyJWT not installed. Install with: pip install PyJWT>=2.4.0")
+    
+    # TODO: Implement JWKS signature verification (future work)
+    # For now, perform unverified decode
+    if verify:
+        warnings.warn(
+            "JWT signature verification requested but not yet implemented. "
+            "Using unverified decode. This should not be used in production "
+            "for security-critical operations.",
+            UserWarning,
+            stacklevel=2
+        )
+    
+    try:
+        # Unverified decode - JWKS verification is future work
+        claims = jwt.decode(token, options={"verify_signature": False})
+        
+        # Derive principal from claims
+        # Priority: cognito:username > email > sub
+        principal = (
+            claims.get('cognito:username') or 
+            claims.get('email') or 
+            claims.get('sub')
+        )
+        
+        if principal:
+            claims['principal'] = principal
+        
+        return claims
+    
+    except jwt.DecodeError as e:
+        raise ValueError(f"Invalid JWT token: {e}")
+    except Exception as e:
+        raise ValueError(f"Failed to decode JWT token: {e}")
+
+
+def derive_principal(claims: Dict[str, Any]) -> str:
+    """
+    Derive a principal identifier from identity claims.
+    
+    Args:
+        claims: Identity claims dictionary
+    
+    Returns:
+        str: Principal identifier
+    
+    Raises:
+        ValueError: If no suitable principal identifier found
+    """
+    principal = (
+        claims.get('principal') or
+        claims.get('cognito:username') or
+        claims.get('email') or
+        claims.get('sub')
+    )
+    
+    if not principal:
+        raise ValueError("No principal identifier found in claims")
+    
+    return str(principal)
+
+
+def is_admin(claims: Dict[str, Any]) -> bool:
+    """
+    Check if the identity has admin privileges.
+    
+    Args:
+        claims: Identity claims dictionary
+    
+    Returns:
+        bool: True if user has admin privileges
+    
+    Note:
+        Currently checks for 'cognito:groups' containing 'admin'.
+        Extend this logic as needed for your authorization model.
+    """
+    groups = claims.get('cognito:groups', [])
+    if isinstance(groups, list):
+        return 'admin' in groups
+    return False

aimodelshare/aws.py

@@ -0,0 +1,511 @@
+import os
+import boto3
+import botocore
+import requests
+import json
+import base64
+from aimodelshare.exceptions import AuthorizationError, AWSAccessError
+from aimodelshare.modeluser import get_jwt_token
+
+def set_credentials(credential_file=None, type="submit_model", apiurl="apiurl", manual = True, cloud="aws"):
+  import os
+  import getpass
+  from aimodelshare.aws import get_aws_token
+  from aimodelshare.modeluser import get_jwt_token, setup_bucket_only
+  if all([credential_file==None, type=="submit_model"]):
+    set_credentials_public(type="submit_model", apiurl=apiurl)
+    os.environ["AWS_TOKEN"]=get_aws_token()
+  elif all([credential_file==None, type=="deploy_model", cloud=="model_share"]):
+    set_credentials_public_aimscloud(credential_file=None, type="deploy_model", apiurl=apiurl)
+    os.environ["cloud_location"] = cloud
+  else:
+      ##TODO: Require that "type" is provided, to ensure correct env vars get loaded
+      flag = False
+
+      # Set AI Modelshare Username & Password
+      if all([manual == True, credential_file==None]):
+        user = getpass.getpass(prompt="Modelshare.ai Username:")
+        os.environ["username"] = user
+        pw = getpass.getpass(prompt="Modelshare.ai Password:")
+        os.environ["password"] = pw
+        
+      else: 
+        f = open(credential_file)
+
+        for line in f:
+          if "aimodelshare_creds" in line or "AIMODELSHARE_CREDS" in line:
+            for line in f:
+              if line == "\n":
+                break
+              try:
+                value = line.split("=", 1)[1].strip()
+                value = value[1:-1]
+                key = line.split("=", 1)[0].strip()
+                os.environ[key.lower()] = value
+
+              except LookupError: 
+                print(* "Warning: Review format of", credential_file, ". Format should be variablename = 'variable_value'")
+                break
+      
+      #Validate Username & Password
+      try: 
+        os.environ["AWS_TOKEN"]=get_aws_token()
+
+        print("Modelshare.ai login credentials set successfully.")
+      except: 
+        print("Credential confirmation unsuccessful. Check username & password and try again.")
+        return
+      
+      # Set AWS Creds Manually (submit or deploy)
+      if  all([manual == True,credential_file==None]):
+        flag = True
+        access_key = getpass.getpass(prompt="AWS_ACCESS_KEY_ID:")
+        os.environ["AWS_ACCESS_KEY_ID"] = access_key
+        os.environ["AWS_ACCESS_KEY_ID_AIMS"] = access_key
+
+        secret_key = getpass.getpass(prompt="AWS_SECRET_ACCESS_KEY:")
+        os.environ["AWS_SECRET_ACCESS_KEY"] = secret_key
+        os.environ["AWS_SECRET_ACCESS_KEY_AIMS"] = secret_key
+
+        region = getpass.getpass(prompt="AWS_REGION:")
+        os.environ["AWS_REGION"] = region
+        os.environ["AWS_REGION_AIMS"] = region
+
+      # Set AWS creds from file
+      else:  
+        f = open(credential_file)
+        if type == "submit_model": 
+          for line in f:
+            if (apiurl in line) and ((type in line) or (type.upper() in line)): ## searches on apiurl AND type
+              flag = True
+              for line in f:
+                if line == "\n":
+                  break
+                try:
+                  value = line.split("=", 1)[1].strip()
+                  value = value[1:-1]
+                  key = line.split("=", 1)[0].strip()
+                  os.environ[key.upper()] = value
+                except LookupError: 
+                  print(* "Warning: Review format of", credential_file, ". Format should be variablename = 'variable_value'.")
+                  break
+
+        elif type == "deploy_model": 
+          for line in f:
+            if ((type in line) or (type.upper() in line)):  ## only searches on type
+              flag = True
+              for line in f:
+                if line == "\n":
+                  break
+                try:
+                  value = line.split("=", 1)[1].strip()
+                  value = value[1:-1]
+                  key = line.split("=", 1)[0].strip()
+                  os.environ[key.upper()] = value
+                except LookupError: 
+                  print(* "Warning: Review format of", credential_file, ". Format should be variablename = 'variable_value'.")
+                  break
+
+      if "AWS_ACCESS_KEY_ID_AIMS" in os.environ:
+        pass
+      elif "AWS_ACCESS_KEY_ID" in os.environ:
+          os.environ['AWS_ACCESS_KEY_ID_AIMS']=os.environ.get("AWS_ACCESS_KEY_ID")
+      if "AWS_SECRET_ACCESS_KEY_AIMS" in os.environ:
+        pass
+      elif "AWS_SECRET_ACCESS_KEY" in os.environ:
+          os.environ['AWS_SECRET_ACCESS_KEY_AIMS']=os.environ.get("AWS_SECRET_ACCESS_KEY")
+
+      if 'AWS_REGION_AIMS' in os.environ:
+        pass
+      elif "AWS_REGION" in os.environ:
+          os.environ['AWS_REGION_AIMS']=os.environ.get("AWS_REGION")
+      # Validate AWS Creds 
+      import boto3
+      try: 
+        client = boto3.client('sts', aws_access_key_id=os.environ.get("AWS_ACCESS_KEY_ID"), aws_secret_access_key=os.environ.get("AWS_SECRET_ACCESS_KEY_AIMS"))
+        details = client.get_caller_identity()
+        print("AWS credentials set successfully.")
+      except: 
+        print("AWS credential confirmation unsuccessful. Check AWS_ACCESS_KEY_ID & AWS_SECRET_ACCESS_KEY and try again.")
+        return
+      
+      # Set Environment Variables for deploy models
+      if type == "deploy_model":
+        get_jwt_token(os.environ.get("username"), os.environ.get("password"))
+        setup_bucket_only()  # Use new function that doesn't create IAM users  
+        
+      if not flag: 
+        print("Error: apiurl or type not found in"+str(credential_file)+". Please correct entries and resubmit.")
+      
+      try:
+        f.close()
+      except:
+        pass
+
+  return
+
+def set_credentials_public(credential_file=None, type="submit_model", apiurl="apiurl", manual = True):
+  import os
+  import getpass
+  from aimodelshare.aws import get_aws_token
+  from aimodelshare.modeluser import get_jwt_token, setup_bucket_only
+
+  ##TODO: Require that "type" is provided, to ensure correct env vars get loaded
+  flag = False
+
+  # Set AI Modelshare Username & Password
+  if all([manual == True, credential_file==None]):
+    user = getpass.getpass(prompt="Modelshare.ai Username:")
+    os.environ["username"] = user
+    pw = getpass.getpass(prompt="Modelshare.ai Password:")
+    os.environ["password"] = pw
+  
+  else: 
+    f = open(credential_file)
+
+    for line in f:
+      if "aimodelshare_creds" in line or "AIMODELSHARE_CREDS" in line:
+        for line in f:
+          if line == "\n":
+            break
+          try:
+            value = line.split("=", 1)[1].strip()
+            value = value[1:-1]
+            key = line.split("=", 1)[0].strip()
+            os.environ[key.lower()] = value
+
+          except LookupError: 
+            print(* "Warning: Review format of", credential_file, ". Format should be variablename = 'variable_value'")
+            break
+    if "AWS_ACCESS_KEY_ID_AIMS" in os.environ:
+      pass
+    elif "AWS_ACCESS_KEY_ID" in os.environ:
+        os.environ['AWS_ACCESS_KEY_ID_AIMS']=os.environ.get("AWS_ACCESS_KEY_ID")
+    if "AWS_SECRET_ACCESS_KEY_AIMS" in os.environ:
+      pass
+    elif "AWS_SECRET_ACCESS_KEY" in os.environ:
+        os.environ['AWS_SECRET_ACCESS_KEY_AIMS']=os.environ.get("AWS_SECRET_ACCESS_KEY")
+
+    if 'AWS_REGION_AIMS' in os.environ:
+      pass
+    elif "AWS_REGION" in os.environ:
+        os.environ['AWS_REGION_AIMS']=os.environ.get("AWS_REGION")
+  #Validate Username & Password
+  try: 
+    os.environ["AWS_TOKEN"]=get_aws_token()
+
+    print("Modelshare.ai login credentials set successfully.")
+  except: 
+    print("Credential confirmation unsuccessful. Check username & password and try again.")
+    return
+  
+   # Set AWS creds from file
+ 
+  try:
+    f.close()
+  except:
+    pass
+
+  return
+
+def set_credentials_public_aimscloud(credential_file=None, type="deploy_model", apiurl="apiurl", manual = True):
+  import os
+  import getpass
+  from aimodelshare.aws import get_aws_token
+  from aimodelshare.modeluser import get_jwt_token, setup_bucket_only
+
+  ##TODO: Require that "type" is provided, to ensure correct env vars get loaded
+  flag = False
+
+  # Set AI Modelshare Username & Password
+  if all([manual == True, credential_file==None]):
+    user = getpass.getpass(prompt="Modelshare.ai Username:")
+    os.environ["username"] = user
+    pw = getpass.getpass(prompt="Modelshare.ai Password:")
+    os.environ["password"] = pw
+  
+  else: 
+    f = open(credential_file)
+
+    for line in f:
+      if "aimodelshare_creds" in line or "AIMODELSHARE_CREDS" in line:
+        for line in f:
+          if line == "\n":
+            break
+          try:
+            value = line.split("=", 1)[1].strip()
+            value = value[1:-1]
+            key = line.split("=", 1)[0].strip()
+            os.environ[key.lower()] = value
+
+          except LookupError: 
+            print(* "Warning: Review format of", credential_file, ". Format should be variablename = 'variable_value'")
+            break
+    if "AWS_ACCESS_KEY_ID_AIMS" in os.environ:
+      pass
+    elif "AWS_ACCESS_KEY_ID" in os.environ:
+        os.environ['AWS_ACCESS_KEY_ID_AIMS']=os.environ.get("AWS_ACCESS_KEY_ID")
+    if "AWS_SECRET_ACCESS_KEY_AIMS" in os.environ:
+      pass
+    elif "AWS_SECRET_ACCESS_KEY" in os.environ:
+        os.environ['AWS_SECRET_ACCESS_KEY_AIMS']=os.environ.get("AWS_SECRET_ACCESS_KEY")
+
+    if 'AWS_REGION_AIMS' in os.environ:
+      pass
+    elif "AWS_REGION" in os.environ:
+        os.environ['AWS_REGION_AIMS']=os.environ.get("AWS_REGION")
+  #Validate Username & Password
+  try: 
+    os.environ["AWS_TOKEN"]=get_aws_token()
+    get_jwt_token(user,pw)
+    print("Modelshare.ai login credentials set successfully.")
+  except: 
+    print("Credential confirmation unsuccessful. Check username & password and try again.")
+    return
+  
+   # Set AWS creds from file
+ 
+  try:
+    f.close()
+  except:
+    pass
+
+  return
+
+def get_aws_token():
+    config = botocore.config.Config(signature_version=botocore.UNSIGNED)
+
+    provider_client = boto3.client(
+        "cognito-idp", region_name="us-east-2", config=config
+    )
+
+    try:
+        response = provider_client.initiate_auth(
+            ClientId="7ptv9f8pt36elmg0e4v9v7jo9t",
+            AuthFlow="USER_PASSWORD_AUTH",
+             AuthParameters={"USERNAME": os.getenv('username'), "PASSWORD": os.getenv('password')},
+        )
+
+    except Exception as err:
+        raise AuthorizationError("Could not authorize user. " + str(err))
+
+    return response["AuthenticationResult"]["IdToken"]
+
+
+def get_token_from_session(session_id):
+    """
+    Retrieve AWS JWT token from session API endpoint.
+    
+    Args:
+        session_id: The session identifier from URL parameter
+        
+    Returns:
+        str: JWT token if session is valid
+        
+    Raises:
+        AuthorizationError: If session is invalid or API request fails
+    """
+    try:
+        # NOTE: This URL should be configured via environment variable or config file
+        # Update AIMODELSHARE_SESSION_API_URL environment variable to override
+        session_api_url = os.getenv(
+            "AIMODELSHARE_SESSION_API_URL",
+            f"https://b22q73wp50.execute-api.us-east-1.amazonaws.com/dev/sessions/{session_id}"
+        )
+        
+        response = requests.get(session_api_url, timeout=10)
+        response.raise_for_status()
+        
+        data = response.json()
+        token = data.get("token") or data.get("id_token") or data.get("IdToken")
+        
+        if not token:
+            raise AuthorizationError("No token found in session API response")
+            
+        return token
+        
+    except requests.RequestException as err:
+        raise AuthorizationError(f"Failed to retrieve token from session: {err}")
+    except (KeyError, ValueError) as err:
+        raise AuthorizationError(f"Invalid session API response: {err}")
+
+
+def _get_username_from_token(token):
+    """
+    Extract username from JWT token claims.
+    
+    Args:
+        token: JWT token string
+        
+    Returns:
+        str: Username extracted from 'cognito:username' or 'username' claim, or None if not found
+    """
+    try:
+        # JWT tokens have 3 parts: header.payload.signature
+        parts = token.split('.')
+        if len(parts) != 3:
+            return None
+            
+        # Decode the payload (second part)
+        # Add padding if needed for base64 decoding
+        payload = parts[1]
+        padding = 4 - (len(payload) % 4)
+        if padding != 4:
+            payload += '=' * padding
+            
+        decoded = base64.urlsafe_b64decode(payload)
+        claims = json.loads(decoded)
+        
+        # Try different possible username claim names
+        # Note: 'sub' is excluded as it's typically a UUID, not a human-readable username
+        username = claims.get('cognito:username') or claims.get('username')
+        
+        return username
+        
+    except (ValueError, json.JSONDecodeError, KeyError, IndexError) as err:
+        print(f"Warning: Could not extract username from token: {err}")
+        return None
+
+
+def get_aws_session(aws_key=None, aws_secret=None, aws_region=None):
+    session = boto3.Session(
+        aws_access_key_id=aws_key,
+        aws_secret_access_key=aws_secret,
+        region_name=aws_region
+    )
+
+def get_aws_client(aws_key=None, aws_secret=None, aws_region=None):
+    key = aws_key if aws_key is not None else os.environ.get("AWS_ACCESS_KEY_ID_AIMS")
+    secret = (
+        aws_secret
+        if aws_secret is not None
+        else os.environ.get("AWS_SECRET_ACCESS_KEY_AIMS")
+    )
+    region = (
+        aws_region if aws_region is not None else os.environ.get("AWS_REGION_AIMS") #changed
+    )
+
+    if any([key is None, secret is None, region is None]):
+        raise ValueError("Invalid arguments")
+
+    usersession = boto3.session.Session(
+        aws_access_key_id=key, aws_secret_access_key=secret, region_name=region,
+    )
+
+    s3client = usersession.client("s3")
+    s3resource = usersession.resource("s3")
+
+    return {"client": s3client, "resource": s3resource}
+
+
+def get_s3_iam_client(aws_key=None,aws_password=None, aws_region=None):
+
+  key = aws_key if aws_key is not None else os.environ.get("AWS_ACCESS_KEY_ID_AIMS")
+  password = (
+        aws_password
+        if aws_password is not None
+        else os.environ.get("AWS_SECRET_ACCESS_KEY_AIMS"))
+  region = (
+        aws_region if aws_region is not None else os.environ.get("AWS_REGION_AIMS")) #changed
+
+  if any([key is None, password is None, region is None]):
+        raise AuthorizationError("Please set your aws credentials before creating your prediction API.")
+
+  usersession = boto3.session.Session(
+        aws_access_key_id=key, aws_secret_access_key=password, region_name=region,
+    )
+
+  s3_client = boto3.client('s3' , region_name ="us-east-1")
+  s3_resource = boto3.resource('s3')
+  iam_client = boto3.client('iam')
+  iam_resource = boto3.resource('iam')
+
+  s3 = {"client":s3_client,"resource":s3_resource}
+  iam = {"client":iam_client,"resource":iam_resource}
+
+  return s3,iam,region
+
+
+def run_function_on_lambda(url, username=None, token=None, **kwargs):
+    if username==None:
+      kwargs["apideveloper"] = os.environ.get("username")
+    else:
+      kwargs["apideveloper"] = username
+    kwargs["apiurl"] = url
+    if token==None:
+      headers_with_authentication = {
+          "content-type": "application/json",
+          "authorizationToken": os.environ.get("AWS_TOKEN"),
+      }
+    else:
+      headers_with_authentication = {
+          "content-type": "application/json",
+          "authorizationToken": token,
+      }
+
+    response = requests.post(
+        "https://bhrdesksak.execute-api.us-east-1.amazonaws.com/dev/modeldata",
+        json=kwargs,
+        headers=headers_with_authentication,
+    )
+
+    if response.status_code != 200:
+        return (
+            None,
+            AWSAccessError(
+                "Error:"
+                + "Please make sure your api url and token are correct."
+            ),
+        )
+
+    return response, None
+
+
+def get_token(username, password):
+      #get token for access to prediction lambas or to submit predictions to generate model evaluation metrics
+      tokenstring = '{\"username\": \"$usernamestring\", \"password\": \"$passwordstring\"}'
+      from string import Template
+      t = Template(tokenstring)
+      newdata = t.substitute(
+          usernamestring=username, passwordstring=password)
+      api_url='https://xgwe1d6wai.execute-api.us-east-1.amazonaws.com/dev' 
+      headers={ 'Content-Type':'application/json'}
+      token =requests.post(api_url,headers=headers,data=json.dumps({"action": "login", "request":newdata}))
+      return token.text
+
+def configure_credentials(): 
+    import getpass
+    user = getpass.getpass(prompt="Modelshare.ai Username:")
+    pw = getpass.getpass(prompt="Modelshare.ai Password:")
+    input_AWS_ACCESS_KEY_ID = getpass.getpass(prompt="AWS_ACCESS_KEY_ID:")
+    input_AWS_SECRET_ACCESS_KEY = getpass.getpass(prompt="AWS_SECRET_ACCESS_KEY:")
+    input_AWS_REGION = getpass.getpass(prompt="AWS_REGION:")
+    
+    
+    #Format output text
+    formatted_userpass = ('[aimodelshare_creds] \n'
+                          'username = "' + user + '"\n'
+                          'password = "' + pw + '"\n\n')
+
+    formatted_new_creds = ("#Deploy Credentials \n"
+                    '[deploy_model]\n'
+                    'AWS_ACCESS_KEY_ID = "' + input_AWS_ACCESS_KEY_ID + '"\n'
+                    'AWS_SECRET_ACCESS_KEY = "' + input_AWS_SECRET_ACCESS_KEY +'"\n'
+                    'AWS_REGION = "' + input_AWS_REGION + '"\n')
+    
+    # Generate .txt file with new credentials 
+    f= open("credentials.txt","w+")
+    f.write(formatted_userpass + formatted_new_creds)
+    f.close()
+    
+    return print("Configuration successful. New credentials file saved as 'credentials.txt'")
+
+__all__ = [
+    get_aws_token,
+    get_aws_client,
+    run_function_on_lambda,
+    get_s3_iam_client,
+    set_credentials,
+    configure_credentials,
+    set_credentials_public
+]