AWS-Panel/backend/aws_ops.py

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)