from __future__ import annotations from typing import Any, Dict, List, Optional import boto3 from botocore.config import Config as BotoConfig from .config import settings from .models import AWSCredential, CredentialType def _boto_config() -> BotoConfig: proxies = None if settings.aws_proxy_url: proxies = {"https": settings.aws_proxy_url, "http": settings.aws_proxy_url} return BotoConfig( connect_timeout=settings.aws_timeout, read_timeout=settings.aws_timeout, proxies=proxies, ) def build_session(credential: AWSCredential, region: str): cfg = _boto_config() if credential.credential_type == CredentialType.ACCESS_KEY: session = boto3.Session( aws_access_key_id=credential.access_key_id, aws_secret_access_key=credential.secret_access_key, region_name=region or credential.default_region, ) return session, cfg # Assume role path base_session = boto3.Session( aws_access_key_id=credential.access_key_id, aws_secret_access_key=credential.secret_access_key, region_name=region or credential.default_region, ) sts = base_session.client("sts", config=cfg, region_name=region or credential.default_region) assume_kwargs: Dict[str, Any] = { "RoleArn": credential.role_arn, "RoleSessionName": "ec2-panel", } if credential.external_id: assume_kwargs["ExternalId"] = credential.external_id resp = sts.assume_role(**assume_kwargs) creds = resp["Credentials"] session = boto3.Session( aws_access_key_id=creds["AccessKeyId"], aws_secret_access_key=creds["SecretAccessKey"], aws_session_token=creds["SessionToken"], region_name=region or credential.default_region, ) return session, cfg def describe_instances( credential: AWSCredential, region: str, filters: Optional[List[Dict[str, Any]]] = None, instance_ids: Optional[List[str]] = None, ) -> Dict[str, Any]: session, cfg = build_session(credential, region) client = session.client("ec2", region_name=region or credential.default_region, config=cfg) params: Dict[str, Any] = {} if filters: params["Filters"] = filters if instance_ids: params["InstanceIds"] = instance_ids return client.describe_instances(**params) def describe_instance_status( credential: AWSCredential, region: str, instance_ids: List[str] ) -> Dict[str, Any]: session, cfg = build_session(credential, region) client = session.client("ec2", region_name=region or credential.default_region, config=cfg) return client.describe_instance_status(InstanceIds=instance_ids, IncludeAllInstances=True) def run_instances( credential: AWSCredential, region: str, ami_id: str, instance_type: str, key_name: Optional[str], security_groups: Optional[List[str]], subnet_id: Optional[str], min_count: int = 1, max_count: int = 1, name_tag: Optional[str] = None, ) -> Dict[str, Any]: session, cfg = build_session(credential, region) client = session.client("ec2", region_name=region or credential.default_region, config=cfg) params: Dict[str, Any] = { "ImageId": ami_id, "InstanceType": instance_type, "MinCount": min_count, "MaxCount": max_count, } if key_name: params["KeyName"] = key_name if security_groups: params["SecurityGroupIds"] = security_groups if subnet_id: params["SubnetId"] = subnet_id if name_tag: params["TagSpecifications"] = [ {"ResourceType": "instance", "Tags": [{"Key": "Name", "Value": name_tag}]} ] return client.run_instances(**params) def start_instances(credential: AWSCredential, region: str, instance_ids: List[str]) -> Dict[str, Any]: session, cfg = build_session(credential, region) client = session.client("ec2", region_name=region or credential.default_region, config=cfg) return client.start_instances(InstanceIds=instance_ids) def stop_instances(credential: AWSCredential, region: str, instance_ids: List[str]) -> Dict[str, Any]: session, cfg = build_session(credential, region) client = session.client("ec2", region_name=region or credential.default_region, config=cfg) return client.stop_instances(InstanceIds=instance_ids) def reboot_instances(credential: AWSCredential, region: str, instance_ids: List[str]) -> Dict[str, Any]: session, cfg = build_session(credential, region) client = session.client("ec2", region_name=region or credential.default_region, config=cfg) return client.reboot_instances(InstanceIds=instance_ids) def terminate_instances(credential: AWSCredential, region: str, instance_ids: List[str]) -> Dict[str, Any]: session, cfg = build_session(credential, region) client = session.client("ec2", region_name=region or credential.default_region, config=cfg) return client.terminate_instances(InstanceIds=instance_ids)