한번에 모든 보안 그룹에 규칙 추가하기

#!/usr/bin/env python3

import csv
import boto3
import botocore
import sys

def add_security_group_rules(sg_file, rules_file):
    """
    보안 그룹 추가

    Args:
        sg_file: 보안 그룹 목록 csv 파일
        rules_file: 추가하려는 규칙 csv 파일
    """

    # AWS EC2 클라이언트 초기화
    try:
        ec2 = boto3.client("ec2")
    except Exception as e:
        print(f"AWS 에러 : {e}")
        return False

    # 보안 그룹 목록 읽기
    sg_dict = {}
    try:
        with open(sg_file, "r", encoding="utf-8") as f:
            reader = csv.DictReader(f)
            for row in reader:
                sg_dict[row["GroupName"]] = row["GroupId"]
        print(f"보안 그룹 {len(sg_dict)} 개를 확인")
    except Exception as e:
        print(f"보안 그룹 읽기 실패 : {e}")
        return False

    # 규칙 읽기 추가
    success_count = 0
    error_count = 0

    try:
        with open(rules_file, "r", encoding="utf-8") as f:
            reader = csv.DictReader(f)
            for row in reader:
                sg_name = row["sg_name"]
                direction = row["direction"]
                protocol = row["protocol"]
                from_port = int(row["from_port"]) if row["from_port"] else None
                to_port = int(row["to_port"]) if row["to_port"] else None
                cidr = row["cidr"]
                description = row.get("description", "")

                # 보안 그룹 ID 확인
                if sg_name not in sg_dict:
                    print(f"다음 보안 그룹 이름의 확인 실패 : {sg_name}")
                    error_count += 1
                    continue

                sg_id = sg_dict[sg_name]

                # 규칙 추가
                if add_rule(
                    ec2,
                    sg_id,
                    direction,
                    protocol,
                    from_port,
                    to_port,
                    cidr,
                    description,
                ):
                    print(
                        f"규칙 추가 성공 : {sg_name} ({direction}) {protocol}:{from_port}-{to_port} {cidr}"
                    )
                    success_count += 1
                else:
                    error_count += 1

    except Exception as e:
        print(f"규칙 파일 읽기 실패 : {e}")
        return False

    print(f"\n결과: 성공 {success_count} 개, 에러 {error_count} 개")
    return error_count == 0

def add_rule(ec2, sg_id, direction, protocol, from_port, to_port, cidr, description):
    """
    개별 규칙을 추가
    """
    try:
        # IP
        ip_permission = {
            "IpProtocol": protocol,
            "IpRanges": [{"CidrIp": cidr, "Description": description}],
        }

        # 포트 지정이 있는 경우
        if from_port is not None:
            ip_permission["FromPort"] = from_port
        if to_port is not None:
            ip_permission["ToPort"] = to_port

        # 인바운드 혹은 아웃바운드 규칙 추가
        if direction.lower() == "inbound":
            ec2.authorize_security_group_ingress(
                GroupId=sg_id, IpPermissions=[ip_permission]
            )
        elif direction.lower() == "outbound":
            ec2.authorize_security_group_egress(
                GroupId=sg_id, IpPermissions=[ip_permission]
            )
        else:
            print(f"잘못된 설정 : {direction}")
            return False

        return True

    except botocore.exceptions.ClientError as e:
        error_code = e.response["Error"]["Code"]
        if error_code == "InvalidPermission.Duplicate":
            print(f"다음 그룹에 이미 동일한 규칙이 있습니다 : {sg_id}")
            return True  # 동일한 기존 규칙이 있다면 성공으로 간주
        else:
            print(f"규칙 추가 에러 ({sg_id}): {e}")
            return False
    except Exception as e:
        print(f"이 외의 에러 ({sg_id}): {e}")
        return False

def main():
    sg_file = sys.argv[1]
    rules_file = sys.argv[2]

    print("AWS 보안 그룹에 규칙 추가 개시")
    success = add_security_group_rules(sg_file, rules_file)

    if success:
        print("모든 규칙이 추가 성공")
    else:
        print("일부 규칙이 추가 되지 않았음")
        sys.exit(1)

if __name__ == "__main__":
    main()