updsts 0.8.9__py3-none-any.whl

updsts/__init__.py

@@ -0,0 +1 @@
+# encoding: utf-8-sig

updsts/__main__.py

@@ -0,0 +1,71 @@
+# encoding: utf-8-sig
+
+import os
+import sys
+from pathlib import Path
+from argparse import ArgumentParser
+
+from .cmdparam import *
+from .cmd_handler import *
+from .logutil import get_logger
+
+# ---------------------------------------------------------------------------------------
+def main():
+    argp = ArgumentParser(prog="updsts",
+                          description="Get/Update profile of the aws sts security tokens.")
+    # Register common arguments
+    common = ArgumentParser(add_help=False)
+    common.add_argument(
+        '-v',
+        '--verbose',
+        type=int,
+        default=0,
+        choices=[0, 1, 2],
+        metavar='LEVEL',
+        help='Set the verbosity level (0: normal, 1: verbose, 2: debug).'
+    )
+    user_home = os.path.expanduser("~")
+    common.add_argument(
+        '-c',
+        '--credential_file',
+        type=str,
+        default=None,
+        help='Path to the credetial file where profiles stored. (default: ~/.aws/credentials)'
+    )
+
+    subparsers = argp.add_subparsers(dest='command',
+                                     help='Available commands')
+    # Register subcommands
+    register_sub_get(subparsers, handle_get, parent_parser=common)
+    register_sub_list(subparsers, handle_list, parent_parser=common)
+    register_sub_mcp(subparsers, handle_mcp, parent_parser=common)
+
+    try:
+        # Parse the command line arguments
+        args = argp.parse_args()
+        # If no command is specified, show help
+        if args.command is None:
+            argp.print_help()
+        else:
+            # Initialize the logger with the specified verbosity level
+            _ = get_logger(verbose_level=args.verbose)
+            # Execute the handler for the specified command
+            if hasattr(args, 'handler'):
+                args.handler(args)
+            else:
+                argp.print_help()
+
+    except FileNotFoundError:
+        print("Error: Secrets file not found. Please ensure the file exists or specify a valid path.")
+    except PermissionError:
+        print("Error: Permission denied when accessing the credential file. Check your permissions.")
+    except KeyError as e:
+        print(f"Error: Profile name '{e}' is not found in the credential file. Please check the profile name.")
+    except ValueError as e:
+        print(f"Error: {e}")
+    except Exception as e:
+        print(f"Exception: {e}")
+
+# ---------------------------------------------------------------------------------------
+if __name__ == "__main__":
+    main()

updsts/awsutil.py

@@ -0,0 +1,308 @@
+# encoding: utf-8-sig
+
+import os
+import sys
+import boto3
+from datetime import datetime, timezone
+
+from botocore.exceptions import BotoCoreError, ClientError
+from configparser import ConfigParser, NoSectionError, NoOptionError
+from pathlib import Path
+from typing import Optional, Dict, Any, Callable
+
+from .logutil import get_logger
+from .upcred import CredentialUpdater
+
+# ----------------------------------------------------------------------------
+def mask_string(s: str,
+                unmask_chars_head: int = 4,
+                unmask_chars_tail: int = 4,
+                max_strlen = 16) -> str:
+    """
+    Mask the middle characters of the input string `s`, leaving head and tail visible.
+
+    Args:
+        s (str): The input string to mask
+        unmask_chars_head (int, optional): Number of characters to show at the beginning. Defaults to 4.
+        unmask_chars_tail (int, optional): Number of characters to show at the end. Defaults to 4.
+        max_strlen (int, optional): Maximum display length before truncation. Defaults to 16.
+
+    Returns:
+        str: Masked string with the following behavior:
+            - If string length <= unmask_chars_head + unmask_chars_tail: returns the string as is
+            - If string length > max_strlen: shows head + masked chars + tail + ' (total_length chars)'
+            - Otherwise: shows head + masked middle + tail
+
+    Examples:
+        mask_string("abc", 4, 4) -> "abc"
+        mask_string("abcdefgh", 4, 0) -> "abcd****"
+        mask_string("abcdefgh", 4, 2) -> "abcd**gh"
+        mask_string("very_long_string_example", 4, 4, 16) -> "very********mple (24 chars)"
+    """
+    if not s:
+        return ''
+
+    # Strip whitespace and newlines from input
+    s = s.strip()
+    str_len = len(s)
+    if str_len <= unmask_chars_head + unmask_chars_tail:
+        return s
+
+    tail_part = s[-unmask_chars_tail:] if unmask_chars_tail > 0 else ''
+
+    # Handle long strings
+    if str_len > max_strlen:
+        truncate_len = max_strlen
+        middle_mask_len = truncate_len - unmask_chars_head - unmask_chars_tail
+        if middle_mask_len <= 0:
+            return s[:unmask_chars_head] + tail_part + f' ({str_len} chars)'
+        masked_part = '*' * middle_mask_len
+        return s[:unmask_chars_head] + masked_part + tail_part + f' ({str_len} chars)'
+
+    # Normal masking for strings within max_strlen
+    middle_mask_len = str_len - unmask_chars_head - unmask_chars_tail
+    masked_part = '*' * middle_mask_len
+    return s[:unmask_chars_head] + masked_part + tail_part
+
+
+# ----------------------------------------------------------------------------
+def get_credential_file_path(credential_file: str | None = None) -> Path:
+    """
+    Get the path to the AWS credentials file.
+    Args:
+        credential_file (str | None, optional): Path to the AWS credentials file. 
+            If None, the default location (~/.aws/credentials) is used. Defaults to None.
+    Returns:
+        credential_file_path (Path): Path object of the credential file.
+    """
+    cred_path  = credential_file.strip() if credential_file is not None else credential_file
+    if not cred_path:
+        user_home = os.path.expanduser("~")
+        credential_file_path = Path(user_home) / ".aws" / "credentials"
+    else:
+        credential_file_path = Path(cred_path)
+    logger = get_logger()
+    logger.debug(f"Using credential file: {credential_file_path}")
+    return credential_file_path
+
+# ----------------------------------------------------------------------------
+def get_sts_token(profile_name: str,
+                  totp_token: str,
+                  duration_seconds: int = 3600,
+                  credential_file: str | None = None) -> Optional[Dict[str, Any]]:
+    """
+    Get temporary STS token using MFA.
+    Args:
+        profile_name (str): The profile name in the AWS credentials file.
+        totp_token (str): The TOTP token of the registerd MFA device. 
+        credential_file (str | None, optional): Path to the AWS credentials file. 
+            If None, the default location (~/.aws/credentials) is used. Defaults to None.
+    Returns:
+        Optional[Dict[str, Any]]: A dictionary containing the temporary STS credentials
+            if successful, None otherwise. 
+    Raises:
+        Exception: If there is an error reading the credentials file or obtaining the STS token.
+    """
+    logger = get_logger()
+
+    if totp_token is None or len(totp_token) == 0:
+        raise ValueError("TOTP token is required to obtain STS token.")
+
+    credential_file = get_credential_file_path(credential_file)
+    if not credential_file.exists():
+        raise FileNotFoundError(f"Credential file '{credential_file}' does not exist.")
+    config = ConfigParser()
+    config.read(credential_file)
+    try:
+        access_key = config.get(profile_name, 'aws_access_key_id')
+        secret_key = config.get(profile_name, 'aws_secret_access_key')
+        mfa_arn    = config.get(profile_name, 'mfa_device_arn')
+    except (NoOptionError) as e:
+        raise Exception(f"Profile '{profile_name}' is missing required options: {e}")
+    except (NoSectionError) as e:
+        raise Exception(f"Error reading profile '{profile_name}' from credentials file: {e}")
+
+    logger.debug(f"Using profile '{profile_name}' with access key '{access_key}' and MFA device ARN '{mfa_arn}'")
+
+    try:
+        # Read proxy settings from environment variables
+        proxies = {}
+        http_proxy = os.environ.get('http_proxy') or os.environ.get('HTTP_PROXY')
+        https_proxy = os.environ.get('https_proxy') or os.environ.get('HTTPS_PROXY')
+        if http_proxy:
+            proxies['http'] = http_proxy
+            logger.debug(f"Using HTTP proxy: {http_proxy}")
+        if https_proxy:
+            proxies['https'] = https_proxy
+            logger.debug(f"Using HTTPS proxy: {https_proxy}")
+
+        # Create a session with explicit credentials
+        session = boto3.Session(
+            aws_access_key_id=access_key,
+            aws_secret_access_key=secret_key
+        )
+
+        # Create STS client with proxy configuration if available
+        client_config = {}
+        if proxies:
+            client_config['proxies'] = proxies
+
+        sts_client = session.client('sts', **client_config)
+        response = sts_client.get_session_token(DurationSeconds=duration_seconds,
+                                                SerialNumber=mfa_arn,
+                                                TokenCode=totp_token)
+        credentials = response['Credentials']
+        logger.info(f"Successfully obtained temporary STS credentials for profile '{profile_name}'")
+
+        # Convert expiration time from UTC to local timezone
+        expiration_utc = credentials['Expiration']
+        # AWS returns naive datetime in UTC, so we need to make it timezone-aware
+        if expiration_utc.tzinfo is None:
+            expiration_utc = expiration_utc.replace(tzinfo=timezone.utc)
+        # Convert to local timezone
+        expiration_local = expiration_utc.astimezone()
+
+        return {
+            'AccessKeyId': credentials['AccessKeyId'],
+            'SecretAccessKey': credentials['SecretAccessKey'],
+            'SessionToken': credentials['SessionToken'],
+            'Expiration': expiration_local.isoformat()
+        }
+    except (BotoCoreError, ClientError) as e:
+        logger.error(f"Error obtaining STS token: {e}")
+        return None
+
+# ----------------------------------------------------------------------------
+def get_profile_info(profile_name: str,
+                     credential_file: str | None = None,
+                     secret_mask: bool = False,
+                     ctx: ConfigParser = None) -> Optional[Dict[str, str]]:
+    """
+    Get the credencials from the AWS credentials file for the specified profile.
+    Args:
+        profile_name (str): The profile name in the AWS credentials file.
+        credential_file (str | None, optional): Path to the AWS credentials file. 
+            If None, the default location (~/.aws/credentials) is used. Defaults to None.
+    Returns:
+        Optional[Dict[str, str]]: A dictionary containing secret details
+            if successful, None otherwise.
+    Raises:
+        Exception: If there is an error reading the profile from the credentials file.
+    """
+    logger = get_logger()
+    config = ctx
+    if config is None:
+        credential_file = get_credential_file_path(credential_file)
+        if not credential_file.exists():
+            raise FileNotFoundError(f"Credential file '{credential_file}' does not exist.")
+        config = ConfigParser()
+        config.read(credential_file)
+    item_dic = None
+    try:
+        item_dic = {
+            'aws_access_key_id': '(not defined)',
+            'aws_secret_access_key': '(not defined)',
+            'aws_session_token': '(not defined)',
+            'mfa_device_arn': '(not defined)',
+            'totp_secret_name': '(not defined)',
+            'expiration_datetime': '(not defined)'
+        }
+        for key, value in config.items(profile_name):
+            # Remove whitespace and newlines from log output
+            raw_sensitive_value = value.strip() if value else value
+            logger.debug(f"Profile '{profile_name}': {key} = {raw_sensitive_value}")
+            # Mask secret key and session token if required
+            if key == 'aws_access_key_id':
+                secret_key_id = mask_string(raw_sensitive_value) if (raw_sensitive_value and secret_mask) else raw_sensitive_value
+                item_dic[key] = secret_key_id
+            elif key == 'aws_secret_access_key':
+                secret_key = mask_string(raw_sensitive_value) if (raw_sensitive_value and secret_mask) else raw_sensitive_value
+                item_dic[key] = secret_key
+            elif key == 'aws_session_token':
+                session_token = mask_string(raw_sensitive_value) if (raw_sensitive_value and secret_mask) else raw_sensitive_value
+                item_dic[key] = session_token
+            else:
+                item_dic[key] = raw_sensitive_value
+        item_dic['profile_name'] = profile_name
+        return item_dic
+    except (NoSectionError) as e:
+        logger.error(f"Error reading profile '{profile_name}' from credentials file: {e}")
+    except (NoOptionError) as e:
+        logger.error(f"Profile '{profile_name}' is missing required options: {e}")
+    return item_dic
+
+# ----------------------------------------------------------------------------
+def get_profile_list(credential_file: str | None = None,
+                     secret_mask: bool = False) -> list[dict[str, str]]:
+    """
+    Get the list of credencials from the AWS credentials file.
+    Args:
+        credential_file (str | None, optional): Path to the AWS credentials file. 
+            If None, the default location (~/.aws/credentials) is used. Defaults to None.
+    Returns:
+        Optional[list[Dict[str, str]]]: A list of dictionaries containing secret details
+            if successful, None otherwise.
+    Raises:
+        Exception: If there is an error reading the profiles file.
+    """
+    logger = get_logger()
+    credential_file = get_credential_file_path(credential_file)
+    if not credential_file.exists():
+        raise FileNotFoundError(f"Credential file '{credential_file}' does not exist.")
+    config = ConfigParser()
+    config.read(credential_file)
+    profiles = []
+    for section in config.sections():
+        item_dic = get_profile_info(profile_name=section,
+                                    credential_file=credential_file,
+                                    secret_mask=secret_mask,
+                                    ctx=config)
+        if item_dic:
+            profiles.append(item_dic)
+
+    if not profiles:
+        logger.info("No valid profiles found in the credentials file.")
+    else:
+        logger.info(f"Found {len(profiles)} profiles in the credentials file.")
+    return profiles
+
+# ----------------------------------------------------------------------------
+def update_credentials(profile_name: str,
+                       totp_token: str,
+                       duration: int = 3600,
+                       sts_profile_name: str | None = None,
+                       target_key: str | None = None,
+                       cred_file: str | os.PathLike | None = None,
+                       message_handler: Callable[[str], None] | None = None) -> dict[str, str] | None:
+    """
+    Update the AWS credentials file with new STS tokens.
+    message_handler receives user-facing progress messages (e.g. print for CLI).
+    Leave it None in MCP mode so that nothing is written to stdout.
+    """
+    logger = get_logger()
+    ret = None
+    try:
+        sts_credentials = get_sts_token(profile_name=profile_name,
+                                        totp_token=totp_token,
+                                        credential_file=cred_file,
+                                        duration_seconds=duration)
+        if sts_credentials:
+            target_key = profile_name if target_key is None else target_key
+            credential_file_path = Path(cred_file) if cred_file else get_credential_file_path()
+            updater = CredentialUpdater(credential_path=credential_file_path)
+            updater.set_target_tag_name(target_key)
+            updater.set_credentials(sts_credentials)
+            updater.set_sts_profile_name(sts_profile_name)
+            updater.set_message_handler(message_handler)
+            ret = updater.update_credential_file()
+            # The result may be returned to the LLM via MCP, so never expose the raw access key id.
+            if ret and ret.get('aws_access_key_id'):
+                ret['aws_access_key_id'] = mask_string(ret['aws_access_key_id'])
+            logger.info(f"STS Credentials for profile '{profile_name}' updated successfully.")
+        else:
+            logger.error("Failed to retrieve STS credentials.")
+            raise Exception("Failed to retrieve STS credentials.")
+    except Exception as e:
+        logger.error(f"Error: {e}")
+        raise
+    return ret

updsts/cmd_handler.py

@@ -0,0 +1,81 @@
+# encoding: utf-8-sig
+
+import os
+import sys
+from pathlib import Path
+from argparse import ArgumentParser
+
+from .logutil import get_logger
+from .cmdparam import *
+from .awsutil import *
+from .mcp_server import *
+
+# ----------------------------------------------------------------------------
+def handle_get(args):
+    """
+    Handle the 'get' command to generate a TOTP token for a given secret name.
+
+    Args:
+        args (argparse.Namespace): Parsed command line arguments.
+    """
+    profile_name = args.profile_name
+    totp_token = args.totp_token
+    cred_file = args.credential_file if args.credential_file else None
+    sts_profile_name = args.sts_profile_name if hasattr(args, 'sts_profile_name') and args.sts_profile_name else None
+    duration = args.duration if args.duration else 3600
+    sts_profile_name = args.sts_profile_name if hasattr(args, 'sts_profile_name') and args.sts_profile_name else None
+    target_key = args.target_key if hasattr(args, 'target_key') and args.target_key else None
+
+    ret = update_credentials(profile_name=profile_name,
+                             totp_token=totp_token,
+                             duration=duration,
+                             sts_profile_name=sts_profile_name,
+                             target_key=target_key,
+                             cred_file=cred_file,
+                             message_handler=print)
+    if ret:
+        print(f"STS Credentials of profile '{profile_name}' updated successfully.")
+        print(f"The temporary credential({ret.get('updated_profile_name', '')}) will expire at: {ret.get('aws_token_expiration', '')}")
+
+# ----------------------------------------------------------------------------
+def handle_list(args):
+    """
+    Handle the 'list' command to display all registered secrets.
+
+    Args:
+        args (argparse.Namespace): Parsed command line arguments.
+    """
+    credential_file = args.credential_file if args.credential_file else None
+    profiles = get_profile_list(credential_file=credential_file,
+                                secret_mask=True)
+    if not profiles:
+        print("No profiles found.")
+    else:
+        for prof in profiles:
+            print(f"Profile Name: {prof['profile_name']}")
+            print(f"  Access Key ID       : {prof['aws_access_key_id']}")
+            print(f"  Secret Key          : {prof['aws_secret_access_key']}")
+            print(f"  Session Token       : {prof['aws_session_token']}") 
+            print(f"  Expiration DateTime : {prof['expiration_datetime']}")
+            print(f"  MFA Device ARN      : {prof['mfa_device_arn']}")
+            print(f"  TOTP Secret Name    : {prof['totp_secret_name']}")
+            print()
+
+# ----------------------------------------------------------------------------
+def handle_mcp(args):
+    """
+    Handle the 'mcp' command to run the MCP server.
+
+    Args:
+        args (argparse.Namespace): Parsed command line arguments.
+    """
+
+    run_mcp = args.mcp_server if args.mcp_server else False
+    if run_mcp:
+        # If the MCP server flag is set, run the MCP server
+        run_as_mcp_server()
+        pass
+    else:
+        # Otherwise, run the MCP server test
+        disp_tools()
+        pass

updsts/cmdparam.py

@@ -0,0 +1,90 @@
+# encoding: utf-8-sig
+
+import argparse
+
+# ----------------------------------------------------------------------------
+def register_sub_get(subparsers,
+                     handle_get: callable,
+                     parent_parser: argparse.ArgumentParser):
+    """
+    Register the 'update' subcommand to the argument parser.
+    """
+    get_parser = subparsers.add_parser(
+        'get',
+        help='Get/Update new sts token',
+        description='Get/Update new sts token.',
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+        parents=[parent_parser]
+    )
+    get_parser.add_argument(
+        '-n',
+        '--profile_name',
+        type=str,
+        required=True,
+        help='Profile name to get sts secret key.'
+    )
+    get_parser.add_argument(
+        '-sn',
+        '--sts_profile_name',
+        type=str,
+        required=False,
+        default=None,
+        help='STS Profile name in the AWS credentials file.'
+    )
+    get_parser.add_argument(
+        '-t',
+        '--totp_token',
+        type=str,
+        required=True,
+        help='MFA TOTP token of the user'
+    )
+    get_parser.add_argument(
+        '-d',
+        '--duration',
+        type=int,
+        required=False,
+        default=3600,
+        help='Duration seconds of the sts token (default: 3600)'
+    )
+    get_parser.set_defaults(handler=handle_get)
+    return subparsers
+
+# ----------------------------------------------------------------------------
+def register_sub_list(subparsers,
+                      handle_list: callable,
+                      parent_parser: argparse.ArgumentParser):
+    """
+    Register the 'list' subcommand to the argument parser.
+    """
+    list_parser = subparsers.add_parser(
+        'list',
+        help='List all profile names',
+        description='List all profile names.',
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+        parents=[parent_parser]
+    )
+    # Set the function to handle the 'list' command
+    list_parser.set_defaults(handler=handle_list)
+    return subparsers
+
+# ----------------------------------------------------------------------------
+def register_sub_mcp(subparsers,
+                     handle_mcp: callable,
+                     parent_parser: argparse.ArgumentParser):
+    """
+    Register the 'mcp' subcommand to the argument parser.
+    """
+    mcp_parser = subparsers.add_parser(
+        'mcp',
+        help='Subcommand for mcp functions',
+        description='Manage TOTP secrets using MCP tools.',
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+        parents=[parent_parser]
+    )
+    mcp_parser.add_argument(
+        "--mcp-server",
+        action="store_true",
+        help="Run as MCP server"
+    )
+    mcp_parser.set_defaults(handler=handle_mcp)
+    return subparsers

updsts/logutil.py

@@ -0,0 +1,109 @@
+# encoding: utf-8-sig
+
+import os
+import sys
+import logging
+from pathlib import Path
+from typing import Optional
+
+# globals
+logger_obj: Optional[logging.Logger] = None
+is_initialized: bool = False
+
+DEFAULT_LOGGER_NAME = "awssts"
+DEFALT_LOG_LEVEL = logging.WARNING
+DEFAULT_FILE_LEVEL = logging.WARNING
+DEFAULT_CONSOLE_LEVEL = logging.WARNING
+
+# ----------------------------------------------------------------------------
+def get_with_init(
+    log_file: str | None = None,
+    *,
+    level: int = DEFALT_LOG_LEVEL,
+    file_level: int = DEFAULT_FILE_LEVEL,
+    console_level: int = DEFAULT_CONSOLE_LEVEL,
+    force: bool = False
+) -> logging.Logger:
+    """
+    Initialize the logger with the specified configuration.
+    If the logger is already initialized, it will return the existing logger
+
+    Args:
+        log_file (str | None, optional): Defaults to None.
+        level (int, optional): Defaults to DEFAULT_LEVEL.
+        file_level (int, optional):  Defaults to DEFAULT_FILE_LEVEL.
+        console_level (int, optional):  Defaults to DEFAULT_CONSOLE_LEVEL.
+        force (bool, optional):  Defaults to False.
+
+    Returns:
+        logging.Logger: Logger instance for the mktotp module.
+    """
+
+    if log_file is None:
+        user_home = os.path.expanduser("~")
+        log_dir = Path(user_home) / ".awscm" / "log"
+        os.makedirs(log_dir, exist_ok=True)
+        log_file = log_dir / "awscm.log"
+    else:
+        # make the directory for the specified path if it does not exist
+        try:
+            os.makedirs(os.path.dirname(log_file), exist_ok=True)
+        except OSError:
+            # If the directory cannot be created, we just ignore it
+            pass
+
+    global logger_obj, is_initialized
+    if not is_initialized:
+        logger_obj = logging.getLogger(DEFAULT_LOGGER_NAME)
+        logger_obj.setLevel(level=level)
+        formatter = logging.Formatter(
+            '%(asctime)s [%(name)s] [%(levelname)s] %(message)s'
+        )
+
+        try:
+            file_handler = logging.FileHandler(log_file, encoding="utf-8")
+            file_handler.setFormatter(formatter)
+            file_handler.setLevel(file_level)
+            logger_obj.addHandler(file_handler)
+        except (OSError, PermissionError):
+            # If file handler creation fails, we just skip it and only use console handler
+            pass
+
+        console_handler = logging.StreamHandler(sys.stderr)
+        console_handler.setFormatter(formatter)
+        console_handler.setLevel(console_level)
+        logger_obj.addHandler(console_handler)
+
+        is_initialized = True
+    return logger_obj
+
+# ----------------------------------------------------------------------------
+def get_logger(verbose_level: int | None = None) -> logging.Logger:
+    """
+    Module-level shortcut to get the logger.
+
+    Returns:
+        logging.Logger: Logger instance for the mktotp module.
+    """
+    global logger_obj, is_initialized
+
+    # If logger is already initialized and no verbose_level is specified,
+    # return the existing logger instance
+    if is_initialized and verbose_level is None:
+        return logger_obj
+
+    logObj = None
+    if verbose_level is not None:
+        if verbose_level == 0:
+            logObj = get_with_init(level=logging.WARNING)
+        elif verbose_level == 1:
+            logObj = get_with_init(level=logging.INFO)
+        elif verbose_level == 2:
+            logObj = get_with_init(level=logging.DEBUG)
+    else:
+        # If no verbose level is specified, use the default level
+        logObj = get_with_init()
+    return logObj
+
+
+__all__ = ["get_logger"]