aimodelshare 0.1.32__py3-none-any.whl → 0.1.62__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

@@ -10,7 +10,7 @@ def set_credentials(credential_file=None, type="submit_model", apiurl="apiurl",
   import os
   import getpass
   from aimodelshare.aws import get_aws_token
-  from aimodelshare.modeluser import get_jwt_token, create_user_getkeyandpassword
+  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()
@@ -131,7 +131,7 @@ def set_credentials(credential_file=None, type="submit_model", apiurl="apiurl",
       # Set Environment Variables for deploy models
       if type == "deploy_model":
         get_jwt_token(os.environ.get("username"), os.environ.get("password"))
-        create_user_getkeyandpassword()  
+        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.")
@@ -147,7 +147,7 @@ def set_credentials_public(credential_file=None, type="submit_model", apiurl="ap
   import os
   import getpass
   from aimodelshare.aws import get_aws_token
-  from aimodelshare.modeluser import get_jwt_token, create_user_getkeyandpassword
+  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
@@ -211,7 +211,7 @@ def set_credentials_public_aimscloud(credential_file=None, type="deploy_model",
   import os
   import getpass
   from aimodelshare.aws import get_aws_token
-  from aimodelshare.modeluser import get_jwt_token, create_user_getkeyandpassword
+  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

aimodelshare/base_image.py

@@ -8,7 +8,7 @@ import zipfile
 import importlib.resources as pkg_resources
 from string import Template
 
-def lambda_using_base_image(account_id, region, session, project_name, model_dir, requirements_file_path, apiid, memory_size='3000', timeout='90', python_version='3.10'):
+def lambda_using_base_image(account_id, region, session, project_name, model_dir, requirements_file_path, apiid, memory_size='3000', timeout='90', python_version='3.7'):
 
     codebuild_bucket_name=os.environ.get("BUCKET_NAME") # s3 bucket name to create  #TODO: use same bucket and subfolder we used previously to store this data
                                                                                     #Why?  AWS limits users to 100 total buckets!  Our old code only creates one per user per acct.

aimodelshare/containerisation.py

@@ -27,7 +27,7 @@ def create_bucket(s3_client, bucket_name, region):
                 )
         return response
     
-def deploy_container(account_id, region, session, project_name, model_dir, requirements_file_path, apiid, memory_size='1024', timeout='120', python_version='3.10', pyspark_support=False):
+def deploy_container(account_id, region, session, project_name, model_dir, requirements_file_path, apiid, memory_size='1024', timeout='120', python_version='3.7', pyspark_support=False):
 
     codebuild_bucket_name=os.environ.get("BUCKET_NAME") # s3 bucket name to create  #TODO: use same bucket and subfolder we used previously to store this data
                                                         # Why? AWS limits users to 100 total buckets!  Our old code only creates one per user per acct.

aimodelshare/data_sharing/download_data.py

@@ -34,8 +34,30 @@ def progress_bar(layer_label, nb_traits):
 def get_auth_url(registry): # to do with auth
 	return 'https://' + registry + '/token/' # no aws auth
 
-def get_auth_head(auth_url, registry, repository): # to do with auth
-	return get_auth_head_no_aws_auth(auth_url, registry, repository, 'application/vnd.docker.distribution.manifest.v2+json') # no aws auth
+def get_auth_head(auth_url, registry, repository):
+    # Broaden Accept header to allow manifest list / OCI fallbacks
+    return get_auth_head_no_aws_auth(
+        auth_url,
+        registry,
+        repository,
+        ('application/vnd.docker.distribution.manifest.v2+json,'
+         'application/vnd.docker.distribution.manifest.list.v2+json,'
+         'application/vnd.oci.image.manifest.v1+json')
+    )
+
+def _fetch_concrete_manifest(registry, repository, tag_or_digest, auth_head):
+    """Fetch a concrete image manifest (not a list)."""
+    resp = requests.get(
+        f'https://{registry}/v2/{repository}/manifests/{tag_or_digest}',
+        headers=auth_head,
+        verify=False
+    )
+    if not resp.ok:
+        raise RuntimeError(
+            f"Failed to fetch manifest {tag_or_digest} (status {resp.status_code}): {resp.text[:300]}"
+        )
+    return resp
+
 
 def download_layer(layer, layer_count, tmp_img_dir, blobs_resp):
 
@@ -76,87 +98,115 @@ def download_layer(layer, layer_count, tmp_img_dir, blobs_resp):
 	return layer_id, layer_dir
 
 def pull_image(image_uri):
-
-	image_uri_parts = image_uri.split('/')
-
-	registry = image_uri_parts[0]	
-	image, tag = image_uri_parts[2].split(':')
-	repository = '/'.join([image_uri_parts[1], image])
-
-	auth_url = get_auth_url(registry)
-
-	auth_head = get_auth_head(auth_url, registry, repository)
-
-	resp = requests.get('https://{}/v2/{}/manifests/{}'.format(registry, repository, tag), headers=auth_head, verify=False)
-	
-	config = resp.json()['config']['digest']
-	config_resp = requests.get('https://{}/v2/{}/blobs/{}'.format(registry, repository, config), headers=auth_head, verify=False)
-
-	tmp_img_dir = tempfile.gettempdir() + '/' + 'tmp_{}_{}'.format(image, tag)
-	os.mkdir(tmp_img_dir)
-
-	file = open('{}/{}.json'.format(tmp_img_dir, config[7:]), 'wb')
-	file.write(config_resp.content)
-	file.close()
-
-	content = [{
-		'Config': config[7:] + '.json',
-		'RepoTags': [],
-		'Layers': []
-	}]
-	content[0]['RepoTags'].append(image_uri)
-
-	layer_count=0
-	layers = resp.json()['layers'][6:]
-
-	for layer in layers:
-
-		layer_count += 1
-
-		auth_head = get_auth_head(auth_url, registry, repository) # done to keep from expiring
-		blobs_resp = requests.get('https://{}/v2/{}/blobs/{}'.format(registry, repository, layer['digest']), headers=auth_head, stream=True, verify=False)
-
-		layer_id, layer_dir = download_layer(layer, layer_count, tmp_img_dir, blobs_resp)
-		content[0]['Layers'].append(layer_id + '/layer.tar')
-
-		# Creating json file
-		file = open(layer_dir + '/json', 'w')
-
-		# last layer = config manifest - history - rootfs
-		if layers[-1]['digest'] == layer['digest']:
-			json_obj = json.loads(config_resp.content)
-			del json_obj['history']
-			del json_obj['rootfs']
-		else: # other layers json are empty
-			json_obj = json.loads('{}')
-		
-		json_obj['id'] = layer_id
-		file.write(json.dumps(json_obj))
-		file.close()
-
-	file = open(tmp_img_dir + '/manifest.json', 'w')
-	file.write(json.dumps(content))
-	file.close()
-
-	content = {
-		'/'.join(image_uri_parts[:-1]) + '/' + image : { tag : layer_id }
-	}
-
-	file = open(tmp_img_dir + '/repositories', 'w')
-	file.write(json.dumps(content))
-	file.close()
-
-	# Create image tar and clean tmp folder
-	docker_tar = tempfile.gettempdir() + '/' + '_'.join([repository.replace('/', '_'), tag]) + '.tar'
-	sys.stdout.flush()
-
-	tar = tarfile.open(docker_tar, "w")
-	tar.add(tmp_img_dir, arcname=os.path.sep)
-	tar.close()
-
-	shutil.rmtree(tmp_img_dir, onerror=redo_with_write)
-
-	return docker_tar
+    image_uri_parts = image_uri.split('/')
+    registry = image_uri_parts[0]
+    image, tag = image_uri_parts[2].split(':')
+    repository = '/'.join([image_uri_parts[1], image])
+
+    auth_url = get_auth_url(registry)
+    auth_head = get_auth_head(auth_url, registry, repository)
+
+    # 1. Fetch initial manifest (may be list or concrete)
+    resp = _fetch_concrete_manifest(registry, repository, tag, auth_head)
+    manifest_json = resp.json()
+
+    # 2. Handle manifest list fallback
+    if 'config' not in manifest_json:
+        if 'manifests' in manifest_json:
+            # Choose amd64 if available, else first
+            chosen = None
+            for m in manifest_json['manifests']:
+                arch = (m.get('platform') or {}).get('architecture')
+                if arch in ('amd64', 'x86_64'):
+                    chosen = m
+                    break
+            if chosen is None:
+                chosen = manifest_json['manifests'][0]
+            digest = chosen['digest']
+            # Re-auth to avoid token expiry
+            auth_head = get_auth_head(auth_url, registry, repository)
+            resp = _fetch_concrete_manifest(registry, repository, digest, auth_head)
+            manifest_json = resp.json()
+        else:
+            raise KeyError(
+                f"Manifest does not contain 'config' or 'manifests'. Keys: {list(manifest_json.keys())}"
+            )
+
+    if 'config' not in manifest_json or 'layers' not in manifest_json:
+        raise KeyError(
+            f"Unexpected manifest shape. Keys: {list(manifest_json.keys())}"
+        )
+
+    config = manifest_json['config']['digest']
+    config_resp = requests.get(
+        f'https://{registry}/v2/{repository}/blobs/{config}',
+        headers=auth_head,
+        verify=False
+    )
+    if not config_resp.ok:
+        raise RuntimeError(
+            f"Failed to fetch config blob {config} (status {config_resp.status_code}): {config_resp.text[:300]}"
+        )
+
+    tmp_img_dir = tempfile.gettempdir() + '/' + f'tmp_{image}_{tag}'
+    os.mkdir(tmp_img_dir)
+
+    with open(f'{tmp_img_dir}/{config[7:]}.json', 'wb') as f:
+        f.write(config_resp.content)
+
+    content = [{
+        'Config': config[7:] + '.json',
+        'RepoTags': [image_uri],
+        'Layers': []
+    }]
+
+    layer_count = 0
+    layers = manifest_json['layers']  # removed [6:] slicing
+
+    for layer in layers:
+        layer_count += 1
+        # Refresh auth (avoid expiry)
+        auth_head = get_auth_head(auth_url, registry, repository)
+        blobs_resp = requests.get(
+            f'https://{registry}/v2/{repository}/blobs/{layer["digest"]}',
+            headers=auth_head,
+            stream=True,
+            verify=False
+        )
+        if not blobs_resp.ok:
+            raise RuntimeError(
+                f"Failed to stream layer {layer['digest']} status {blobs_resp.status_code}: {blobs_resp.text[:200]}"
+            )
+        layer_id, layer_dir = download_layer(layer, layer_count, tmp_img_dir, blobs_resp)
+        content[0]['Layers'].append(layer_id + '/layer.tar')
+
+        # Create layer json
+        with open(layer_dir + '/json', 'w') as fjson:
+            if layers[-1]['digest'] == layer['digest']:
+                json_obj = json.loads(config_resp.content)
+                json_obj.pop('history', None)
+                json_obj.pop('rootfs', None)
+            else:
+                json_obj = {}
+            json_obj['id'] = layer_id
+            fjson.write(json.dumps(json_obj))
+
+    with open(tmp_img_dir + '/manifest.json', 'w') as mf:
+        mf.write(json.dumps(content))
+
+    # repositories file
+    repositories_json = {
+        '/'.join(image_uri_parts[:-1]) + '/' + image: {tag: layer_id}
+    }
+    with open(tmp_img_dir + '/repositories', 'w') as rf:
+        rf.write(json.dumps(repositories_json))
+
+    docker_tar = tempfile.gettempdir() + '/' + '_'.join([repository.replace('/', '_'), tag]) + '.tar'
+    tar = tarfile.open(docker_tar, "w")
+    tar.add(tmp_img_dir, arcname=os.path.sep)
+    tar.close()
+    shutil.rmtree(tmp_img_dir, onerror=redo_with_write)
+    return docker_tar
 
 
 def extract_data_from_image(image_name, file_name, location):
@@ -215,7 +265,7 @@ def import_quickstart_data(tutorial, section="modelplayground"):
         existing_folder = 'titanic_competition_data'
         
     if all([tutorial == "cars", section == "modelplayground"]):
-        quickstart_repository = "public.ecr.aws/y2e2a1d6/quickstart_car_sales_competition-repository:latest" 
+        quickstart_repository = "public.ecr.aws/z5w0c9e9/quickstart_car_sales_competition-repository:latest" 
         existing_folder = 'used_car_competition_data'
         
     if all([tutorial == "clickbait", section == "modelplayground"]):
@@ -241,7 +291,7 @@ def import_quickstart_data(tutorial, section="modelplayground"):
         existing_folder = 'dog_competition_data'
 
     if all([tutorial == "imdb", section == "modelplayground"]):
-        quickstart_repository = "public.ecr.aws/y2e2a1d6/imdb_quickstart_materials-repository:latest"   
+        quickstart_repository = "public.ecr.aws/z5w0c9e9/imdb_quickstart_materials-repository:latest"
         existing_folder = 'imdb_competition_data'
 
     download_data(quickstart_repository)
@@ -284,8 +334,13 @@ def import_quickstart_data(tutorial, section="modelplayground"):
             #unpack data 
             X_train = pd.read_csv("imdb_quickstart_materials/X_train.csv").squeeze("columns")
             X_test = pd.read_csv("imdb_quickstart_materials/X_test.csv").squeeze("columns")
-            y_test_labels = pd.read_csv("imdb_quickstart_materials/y_test_labels.csv").squeeze("columns")
-            y_train_labels = pd.read_csv("imdb_quickstart_materials/y_train_labels.csv").squeeze("columns")
+            with open("imdb_quickstart_materials/y_train_labels.json", "r") as f:
+                y_train_labels = json.load(f)
+            with open("imdb_quickstart_materials/y_test_labels.json", "r") as f:
+                y_test_labels = json.load(f)
+            import pandas as pd
+            y_train_labels=pd.Series(y_train_labels)
+            y_test_labels=pd.Series(y_test_labels)
             # example data
             example_data = X_train[50:55]
 

aimodelshare/generatemodelapi.py

@@ -23,7 +23,7 @@ from aimodelshare.aws import get_s3_iam_client, run_function_on_lambda, get_toke
 from aimodelshare.bucketpolicy import _custom_upload_policy
 from aimodelshare.exceptions import AuthorizationError, AWSAccessError, AWSUploadError
 from aimodelshare.api import get_api_json
-from aimodelshare.modeluser import create_user_getkeyandpassword
+from aimodelshare.modeluser import decode_token_unverified
 from aimodelshare.preprocessormodules import upload_preprocessor
 from aimodelshare.model import _get_predictionmodel_key, _extract_model_metadata
 from aimodelshare.data_sharing.share_data import share_data_codebuild
@@ -447,7 +447,8 @@ def model_to_api(model_filepath, model_type, private, categorical, y_train, prep
 
     if all([isinstance(email_list, list)]):
         idtoken = get_aws_token()
-        decoded = jwt.decode(idtoken, options={"verify_signature": False})  # works in PyJWT < v2.0
+        decoded = decode_token_unverified(idtoken)
+
         email=None
         email = decoded['email']
         # Owner has to be the first on the list
@@ -592,9 +593,9 @@ def create_competition(apiurl, data_directory, y_test, eval_metric_filepath=None
     """
     if all([isinstance(email_list, list)]):
         if any([len(email_list)>0, public=="True",public=="TRUE",public==True]):
-              import jwt
               idtoken=get_aws_token()
-              decoded = jwt.decode(idtoken, options={"verify_signature": False})  # works in PyJWT < v2.0
+              decoded = decode_token_unverified(idtoken)
+
               email=decoded['email']
               email_list.append(email)
     else:
@@ -757,9 +758,9 @@ def create_experiment(apiurl, data_directory, y_test, eval_metric_filepath=None,
     """
     if all([isinstance(email_list, list)]):
         if any([len(email_list)>0, public=="True",public=="TRUE",public==True]):
-              import jwt
               idtoken=get_aws_token()
-              decoded = jwt.decode(idtoken, options={"verify_signature": False})  # works in PyJWT < v2.0
+              decoded = decode_token_unverified(idtoken)
+
               email=decoded['email']
               email_list.append(email)
     else: