如何为 AWS WAF 速率限制的 IP 配置固定封禁时长

import boto3, os, time, logging

logger = logging.getLogger()
logger.setLevel(logging.INFO)
SCOPE = os.environ['WAF_SCOPE']
wafv2 = boto3.client('wafv2') if SCOPE == 'REGIONAL' else boto3.client('wafv2', region_name='us-east-1')
dynamodb = boto3.resource('dynamodb')

WEBACL_ARN = os.environ['WEBACL_ARN']
STACK_NAME = os.environ['STACK_NAME']
RULES = [r.strip() for r in os.environ['RATE_LIMIT_RULES'].split(',')]
DURATIONS = [int(d.strip()) for d in os.environ['BLOCK_DURATIONS'].split(',')]

def list_all_ip_sets():
    results, marker = [], None
    while True:
        kwargs = {'Scope': SCOPE}
        if marker: kwargs['NextMarker'] = marker
        resp = wafv2.list_ip_sets(**kwargs)
        results.extend(resp.get('IPSets', []))
        marker = resp.get('NextMarker')
        if not marker: break
    return results

def ip_to_cidr(ip):
    return f'{ip}/128' if ':' in ip else f'{ip}/32'

# 在冷启动时缓存 IPSet ID
rule_config = {}
all_ipsets = list_all_ip_sets()
for idx, rule in enumerate(RULES):
    ipset_name = f'{STACK_NAME}-{rule}-BlockedIPs'
    table_name = f'{STACK_NAME}-{rule}-IPs'
    ipset_id = next((s['Id'] for s in all_ipsets if s['Name'] == ipset_name), None)
    dur = DURATIONS[idx] if idx < len(DURATIONS) else DURATIONS[0]
    rule_config[rule] = {'ipset_id': ipset_id, 'ipset_name': ipset_name, 'table': dynamodb.Table(table_name), 'duration': dur}

def process_once():
    end_time = int(time.time())
    start_time = end_time - 15
    for rule, cfg in rule_config.items():
        try:
            resp = wafv2.get_sampled_requests(
                WebAclArn=WEBACL_ARN, RuleMetricName=rule, Scope=SCOPE,
                TimeWindow={'StartTime': start_time, 'EndTime': end_time}, MaxItems=500)
            blocked = {s['Request']['ClientIP'] for s in resp.get('SampledRequests', []) if s.get('Action') == 'BLOCK'}
            if not blocked: continue
            # 先写 DynamoDB 再更新 IPSet
            expire = int(time.time()) + cfg['duration'] * 60
            new_cidrs = set()
            for ip in blocked:
                cidr = ip_to_cidr(ip)
                cfg['table'].put_item(Item={'ip': cidr, 'expire_time': expire, 'ttl': expire})
                new_cidrs.add(cidr)
            # 完整 read-modify-write 重试，每次重试都基于最新 IPSet 状态
            for attempt in range(5):
                try:
                    ipset = wafv2.get_ip_set(Scope=SCOPE, Id=cfg['ipset_id'], Name=cfg['ipset_name'])
                    current = set(ipset['IPSet']['Addresses'])
                    updated = current | new_cidrs
                    if updated != current:
                        wafv2.update_ip_set(Scope=SCOPE, Id=cfg['ipset_id'], Name=cfg['ipset_name'],
                                            Addresses=list(updated), LockToken=ipset['LockToken'])
                        logger.info(f'Added {len(updated - current)} IPs for {rule}')
                    break
                except wafv2.exceptions.WAFOptimisticLockException:
                    if attempt < 4: time.sleep(2 ** attempt)
                    else: raise
        except Exception as e:
            logger.error(f'Error processing {rule}: {e}')

def handler(event, context):
    deadline = time.time() + 55
    while time.time() < deadline:
        process_once()
        remaining = deadline - time.time()
        if remaining > 1:
            time.sleep(1)
        else:
            break

import boto3, json, gzip, os, time, logging
from urllib.parse import unquote_plus

logger = logging.getLogger()
logger.setLevel(logging.INFO)
SCOPE = os.environ['WAF_SCOPE']
s3 = boto3.client('s3')
wafv2 = boto3.client('wafv2') if SCOPE == 'REGIONAL' else boto3.client('wafv2', region_name='us-east-1')
dynamodb = boto3.resource('dynamodb')

STACK_NAME = os.environ['STACK_NAME']
RULES = [r.strip() for r in os.environ['RATE_LIMIT_RULES'].split(',')]
DURATIONS = [int(d.strip()) for d in os.environ['BLOCK_DURATIONS'].split(',')]

def ip_to_cidr(ip):
    return f'{ip}/128' if ':' in ip else f'{ip}/32'

def list_all_ip_sets():
    results, marker = [], None
    while True:
        kwargs = {'Scope': SCOPE}
        if marker: kwargs['NextMarker'] = marker
        resp = wafv2.list_ip_sets(**kwargs)
        results.extend(resp.get('IPSets', []))
        marker = resp.get('NextMarker')
        if not marker: break
    return results

def handler(event, context):
    # 收集所有需要封禁的 IP，按规则分组
    new_ips = {}
    for record in event['Records']:
        bucket = record['s3']['bucket']['name']
        key = unquote_plus(record['s3']['object']['key'])
        obj = s3.get_object(Bucket=bucket, Key=key)
        with gzip.GzipFile(fileobj=obj['Body']) as f:
            for line in f:
                try:
                    log = json.loads(line)
                    if log.get('action') == 'BLOCK':
                        rule_name = log.get('terminatingRuleId', '')
                        if rule_name in RULES:
                            ip = log.get('httpRequest', {}).get('clientIp')
                            if ip:
                                new_ips.setdefault(rule_name, set()).add(ip)
                except: pass

    if not new_ips: return	

    all_ipsets = list_all_ip_sets()
    for rule_name, ips in new_ips.items():
        idx = RULES.index(rule_name)
        dur = DURATIONS[idx] if idx < len(DURATIONS) else DURATIONS[0]
        ipset_name = f'{STACK_NAME}-{rule_name}-BlockedIPs'
        table = dynamodb.Table(f'{STACK_NAME}-{rule_name}-IPs')
        ipset_id = next((s['Id'] for s in all_ipsets if s['Name'] == ipset_name), None)
        if not ipset_id: continue

        # 先写 DynamoDB 再更新 IPSet
        expire = int(time.time()) + dur * 60
        new_cidrs = set()
        for ip in ips:
            cidr = ip_to_cidr(ip)
            # 跳过 DynamoDB 中已存在的记录（避免重置过期计时器）
            if not table.get_item(Key={'ip': cidr}).get('Item'):
                table.put_item(Item={'ip': cidr, 'expire_time': expire, 'ttl': expire})
                new_cidrs.add(cidr)
        if not new_cidrs: continue

        # 完整 read-modify-write 重试，每次重试都基于最新 IPSet 状态
        for attempt in range(5):
            try:
                ipset = wafv2.get_ip_set(Scope=SCOPE, Id=ipset_id, Name=ipset_name)
                current = set(ipset['IPSet']['Addresses'])
                updated = current | new_cidrs
                if updated != current:
                    wafv2.update_ip_set(Scope=SCOPE, Id=ipset_id, Name=ipset_name,
                                        Addresses=list(updated), LockToken=ipset['LockToken'])
                    logger.info(f'Added {len(updated - current)} IPs for {rule_name}')
                break
            except wafv2.exceptions.WAFOptimisticLockException:
                if attempt < 4: time.sleep(2 ** attempt)
                else: raise

import boto3, os, time, logging

logger = logging.getLogger()
logger.setLevel(logging.INFO)
SCOPE = os.environ['WAF_SCOPE']
wafv2 = boto3.client('wafv2') if SCOPE == 'REGIONAL' else boto3.client('wafv2', region_name='us-east-1')
dynamodb = boto3.resource('dynamodb')
STACK_NAME = os.environ['STACK_NAME']
RULES = [r.strip() for r in os.environ['RATE_LIMIT_RULES'].split(',')]

def list_all_ip_sets():
    results, marker = [], None
    while True:
        kwargs = {'Scope': SCOPE}
        if marker: kwargs['NextMarker'] = marker
        resp = wafv2.list_ip_sets(**kwargs)
        results.extend(resp.get('IPSets', []))
        marker = resp.get('NextMarker')
        if not marker: break
    return results

def scan_all_items(table):
    items, last_key = [], None
    while True:
        kwargs = {}
        if last_key: kwargs['ExclusiveStartKey'] = last_key
        resp = table.scan(**kwargs)
        items.extend(resp.get('Items', []))
        last_key = resp.get('LastEvaluatedKey')
        if not last_key: break
    return items

def handler(event, context):
    now = int(time.time())
    all_ipsets = list_all_ip_sets()
    for rule in RULES:
        ipset_name = f'{STACK_NAME}-{rule}-BlockedIPs'
        table = dynamodb.Table(f'{STACK_NAME}-{rule}-IPs')
        items = scan_all_items(table)
        expired = [i['ip'] for i in items if i.get('expire_time', 0) <= now]
        if not expired: continue
        ipset_id = next((s['Id'] for s in all_ipsets if s['Name'] == ipset_name), None)
        if not ipset_id: continue
        # 先更新 IPSet，成功后再删 DynamoDB（避免 DynamoDB 已删但 IPSet 更新失败导致永久封禁）
        for attempt in range(5):
            try:
                ipset = wafv2.get_ip_set(Scope=SCOPE, Id=ipset_id, Name=ipset_name)
                current = set(ipset['IPSet']['Addresses'])
                for ip in expired: current.discard(ip)
                wafv2.update_ip_set(Scope=SCOPE, Id=ipset_id, Name=ipset_name,
                                    Addresses=list(current), LockToken=ipset['LockToken'])
                break
            except wafv2.exceptions.WAFOptimisticLockException:
                if attempt < 4: time.sleep(2 ** attempt)
                else: raise
        for ip in expired:
            table.delete_item(Key={'ip': ip})
        logger.info(f'Removed {len(expired)} expired IPs for {rule}')