main.py

import asyncio
import errno  # Added for e.code == errno.ENOENT
import os
import sys
import time
import traceback

from pyroute2 import (
    AsyncIPRoute,  # IPRoute を AsyncIPRoute に変更
    NetlinkError,
    netns,
)
from pyroute2.netlink.rtnl import TC_H_ROOT

# --- 設定値 ---
PRIV_BR = "br-priv"
PUB_BR = "br-pub"
PRIV_NET = "10.0.1"
PUB_NET = "10.0.2"
SUBNET_MASK = "/24"
NUM_SERVERS = 3
PUB_NS = "ns-pub"
VETH_PUB_NS = "veth-ns-pub"
VETH_PUB_BR = "veth-br-pub"


# --- ユーティリティ関数 ---
def namespace_exists(ns_name):
    """名前空間が存在するかチェックする"""
    try:
        # /var/run/netns ディレクトリが存在しない場合、listnetns() はエラーを出す可能性がある
        # pyroute2 の listnetns() は内部で os.listdir('/var/run/netns') を呼び出すため、
        # FileNotFoundError をキャッチする。
        existing_namespaces = netns.listnetns()
        return ns_name in existing_namespaces
    except FileNotFoundError:
        # /var/run/netns が存在しない場合は、どの名前空間も存在しないとみなす
        print(
            f"     Info: Namespace directory /var/run/netns not found. Assuming namespace {ns_name} does not exist.",
            file=sys.stderr,
        )
        return False
    except Exception as e:
        # その他のエラーが発生した場合 (例: パーミッションエラー)
        print(
            f"     Warning: Error checking namespace existence for {ns_name} via list_ns(): {e}",
            file=sys.stderr,
        )
        # エラー時は存在しないと仮定し、作成を試みるようにする
        return False


def create_namespace(ns_name):
    """名前空間を作成する"""
    if not namespace_exists(ns_name):
        try:
            netns.create(ns_name)  # netns.create を使用
            print(f"     Namespace {ns_name} created via netns.create().")
        except NetlinkError as e:
            if e.code == errno.EEXIST:
                # namespace_exists() のチェックと netns.create() の呼び出しの間に作成された場合
                print(
                    f"     Namespace {ns_name} already exists (EEXIST from netns.create(), likely race)."
                )
            else:
                print(
                    f"     Error creating namespace {ns_name} with netns.create(): {e}",
                    file=sys.stderr,
                )
                raise
    else:
        print(f"     Namespace {ns_name} already exists (pre-checked).")


def remove_namespace(ns_name):
    """名前空間を削除する"""
    if not namespace_exists(ns_name):
        print(f"     Namespace {ns_name} does not exist, no need to remove.")
        return False

    try:
        netns.remove(ns_name)  # netns.remove を使用
        print(f"     Namespace {ns_name} successfully deleted via netns.remove().")
        return True
    except NetlinkError as e_pyroute_netns:
        print(
            f"     netns.remove() for {ns_name} failed: {e_pyroute_netns}. Trying 'ip netns del'.",
            file=sys.stderr,
        )
        return False
    except Exception as e_generic:  # 他の予期せぬエラー
        print(
            f"     Error during removal of namespace {ns_name} with netns.remove(): {e_generic}",
            file=sys.stderr,
        )


async def create_network():
    """
    ネットワーク構成を作成する
    """
    print("Setting up 3-node server private and public network simulation...")

    try:
        async with AsyncIPRoute() as ipr: # async with に変更
            # 1. ブリッジを作成
            print(f"1. Creating and bringing up bridges {PRIV_BR} and {PUB_BR}...")
            for br_name in [PRIV_BR, PUB_BR]:
                # link_lookup を await で呼び出す
                if not await ipr.link_lookup(ifname=br_name):
                    try:
                        # link を await で呼び出す
                        await ipr.link("add", ifname=br_name, kind="bridge")
                        print(f"   Bridge {br_name} created.")
                    except Exception as e:
                        print(
                            f"[ERROR] Failed to create bridge {br_name}: {e}",
                            file=sys.stderr,
                        )
                        traceback.print_exc()
                        raise
                try:
                    # link_lookup を await で呼び出す
                    br_indices = await ipr.link_lookup(ifname=br_name)
                    if br_indices:
                        # link を await で呼び出す
                        await ipr.link("set", index=br_indices[0], state="up")
                    else:
                        raise Exception(
                            f"Bridge {br_name} not found after creation/check for bringing up"
                        )
                except Exception as e:
                    print(
                        f"[ERROR] Failed to bring up bridge {br_name}: {e}",
                        file=sys.stderr,
                    )
                    traceback.print_exc()
                    raise
            print("   Bridges created and are up.")

            # 2. サーバーノード用のネットワーク名前空間を作成
            print(
                f"2. Creating {NUM_SERVERS} server namespaces and connecting them to both bridges..."
            )
            for i in range(1, NUM_SERVERS + 1):
                ns_name = f"ns{i}"
                veth_ns_priv = f"veth-ns{i}-priv"
                veth_br_priv = f"veth-br{i}-priv"
                priv_ip_str = f"{PRIV_NET}.{i}"
                priv_ip_full = f"{priv_ip_str}{SUBNET_MASK}"

                veth_ns_pub = f"veth-ns{i}-pub"
                veth_br_pub = f"veth-br{i}-pub"
                pub_ip_str = f"{PUB_NET}.{i}"
                pub_ip_full = f"{pub_ip_str}{SUBNET_MASK}"

                print(f"   - Creating namespace {ns_name}...")
                try:
                    # create_namespace は同期のまま（必要なら後で修正）
                    create_namespace(ns_name)
                except Exception as e:
                    print(
                        f"[ERROR] Failed to create namespace {ns_name}: {e}",
                        file=sys.stderr,
                    )
                    traceback.print_exc()
                    raise

                print(f"     Waiting for namespace {ns_name} to be ready...")

                # === プライベート通信用vethの設定 ===
                print(f"     Configuring private network for {ns_name}...")
                # link_lookup を await で呼び出す
                if not await ipr.link_lookup(ifname=veth_ns_priv) and not await ipr.link_lookup(
                    ifname=veth_br_priv
                ):
                    try:
                        # link を await で呼び出す
                        await ipr.link(
                            "add", ifname=veth_ns_priv, peer=veth_br_priv, kind="veth"
                        )
                        # link_lookup を await で呼び出す
                        veth_ns_priv_idx = (await ipr.link_lookup(ifname=veth_ns_priv))[0]
                        # link_lookup を await で呼び出す
                        veth_br_priv_idx = (await ipr.link_lookup(ifname=veth_br_priv))[0]
                        # link を await で呼び出す
                        await ipr.link("set", index=veth_ns_priv_idx, state="up")
                        # link を await で呼び出す
                        await ipr.link("set", index=veth_br_priv_idx, state="up")
                    except Exception as e:
                        print(
                            f"[ERROR] Failed to create veth pair {veth_ns_priv}<->{veth_br_priv}: {e}",
                            file=sys.stderr,
                        )
                        traceback.print_exc()
                        raise

                try:
                    print(f"     Moving {veth_ns_priv} to namespace {ns_name}...")
                    # link_lookup を await で呼び出す
                    veth_ns_priv_indices_root = await ipr.link_lookup(ifname=veth_ns_priv)
                    if veth_ns_priv_indices_root:
                        # link を await で呼び出す
                        await ipr.link(
                            "set", index=veth_ns_priv_indices_root[0], net_ns_fd=ns_name
                        )
                        print(
                            f"     Successfully moved {veth_ns_priv} to namespace {ns_name}"
                        )
                    else:
                        # AsyncIPRoute に変更, async with に変更
                        async with AsyncIPRoute(netns=ns_name) as ipr_check_ns:
                             # link_lookup を await で呼び出す
                            if not await ipr_check_ns.link_lookup(ifname=veth_ns_priv):
                                raise Exception(
                                    f"{veth_ns_priv} not found in root or target namespace {ns_name}"
                                )
                            print(
                                f"     {veth_ns_priv} already in namespace {ns_name}."
                            )
                except Exception as e:
                    print(
                        f"[ERROR] Failed to move or verify {veth_ns_priv} in namespace {ns_name}: {e}",
                        file=sys.stderr,
                    )
                    traceback.print_exc()
                    raise

                try:
                    # link_lookup を await で呼び出す
                    veth_br_priv_indices = await ipr.link_lookup(ifname=veth_br_priv)
                    # link_lookup を await で呼び出す
                    priv_br_indices_for_master = await ipr.link_lookup(ifname=PRIV_BR)
                    if veth_br_priv_indices and priv_br_indices_for_master:
                        # link を await で呼び出す
                        await ipr.link(
                            "set",
                            index=veth_br_priv_indices[0],
                            master=priv_br_indices_for_master[0],
                        )
                        # link を await で呼び出す
                        await ipr.link("set", index=veth_br_priv_indices[0], state="up")
                    else:
                        missing = []
                        if not veth_br_priv_indices:
                            missing.append(veth_br_priv)
                        if not priv_br_indices_for_master:
                            missing.append(PRIV_BR)
                        raise Exception(
                            f"Required interface(s) not found for connecting {veth_br_priv} to {PRIV_BR}: {', '.join(missing)}"
                        )
                except Exception as e:
                    print(
                        f"[ERROR] Failed to connect {veth_br_priv} to bridge {PRIV_BR}: {e}",
                        file=sys.stderr,
                    )
                    traceback.print_exc()
                    raise

                try:
                    print(
                        f"     Configuring private IP for {veth_ns_priv} in {ns_name}..."
                    )
                    # AsyncIPRoute に変更, async with に変更
                    async with AsyncIPRoute(netns=ns_name) as ipr_ns:
                         # link_lookup を await で呼び出す
                        veth_indices_in_ns = await ipr_ns.link_lookup(ifname=veth_ns_priv)
                        if veth_indices_in_ns:
                            idx_veth_ns_priv = veth_indices_in_ns[0]
                             # addr を await で呼び出す
                            await ipr_ns.addr(
                                "add",
                                index=idx_veth_ns_priv,
                                address=priv_ip_str,
                                mask=int(SUBNET_MASK[1:]),
                            )
                             # link を await で呼び出す
                            await ipr_ns.link("set", index=idx_veth_ns_priv, state="up")
                            print(
                                f"     Configured private {veth_ns_priv} in {ns_name} with IP {priv_ip_full}"
                            )
                        else:
                            print(
                                f"Warning: Interface {veth_ns_priv} not found in namespace {ns_name} for IP config."
                            )
                except Exception as e:
                    print(
                        f"[ERROR] Failed to configure private IP for {veth_ns_priv} in namespace {ns_name}: {e}",
                        file=sys.stderr,
                    )
                    traceback.print_exc()
                    raise

                # === パブリック通信用vethの設定 === (Similar to private, with pub variables)
                print(f"     Configuring public network for {ns_name}...")
                # link_lookup を await で呼び出す
                if not await ipr.link_lookup(ifname=veth_ns_pub) and not await ipr.link_lookup(
                    ifname=veth_br_pub
                ):
                    try:
                        # link を await で呼び出す
                        await ipr.link(
                            "add", ifname=veth_ns_pub, peer=veth_br_pub, kind="veth"
                        )
                        # link_lookup を await で呼び出す
                        veth_ns_pub_idx = (await ipr.link_lookup(ifname=veth_ns_pub))[0]
                        # link_lookup を await で呼び出す
                        veth_br_pub_idx = (await ipr.link_lookup(ifname=veth_br_pub))[0]
                        # link を await で呼び出す
                        await ipr.link("set", index=veth_ns_pub_idx, state="up")
                        # link を await で呼び出す
                        await ipr.link("set", index=veth_br_pub_idx, state="up")
                    except Exception as e:
                        print(
                            f"[ERROR] Failed to create veth pair {veth_ns_pub}<->{veth_br_pub}: {e}",
                            file=sys.stderr,
                        )
                        traceback.print_exc()
                        raise

                try:
                    print(f"     Moving {veth_ns_pub} to namespace {ns_name}...")
                    # link_lookup を await で呼び出す
                    veth_ns_pub_indices_root = await ipr.link_lookup(ifname=veth_ns_pub)
                    if veth_ns_pub_indices_root:
                        # link を await で呼び出す
                        await ipr.link(
                            "set", index=veth_ns_pub_indices_root[0], net_ns_fd=ns_name
                        )
                        print(
                            f"     Successfully moved {veth_ns_pub} to namespace {ns_name}"
                        )
                    else:
                         # AsyncIPRoute に変更, async with に変更
                        async with AsyncIPRoute(netns=ns_name) as ipr_check_ns:
                             # link_lookup を await で呼び出す
                            if not await ipr_check_ns.link_lookup(ifname=veth_ns_pub):
                                raise Exception(
                                    f"{veth_ns_pub} not found in root or target namespace {ns_name}"
                                )
                            print(f"     {veth_ns_pub} already in namespace {ns_name}.")
                except Exception as e:
                    print(
                        f"[ERROR] Failed to move or verify {veth_ns_pub} in namespace {ns_name}: {e}",
                        file=sys.stderr,
                    )
                    traceback.print_exc()
                    raise

                try:
                    # link_lookup を await で呼び出す
                    veth_br_pub_indices = await ipr.link_lookup(ifname=veth_br_pub)
                    # link_lookup を await で呼び出す
                    pub_br_indices_for_master = await ipr.link_lookup(ifname=PUB_BR)
                    if veth_br_pub_indices and pub_br_indices_for_master:
                        # link を await で呼び出す
                        await ipr.link(
                            "set",
                            index=veth_br_pub_indices[0],
                            master=pub_br_indices_for_master[0],
                        )
                        # link を await で呼び出す
                        await ipr.link("set", index=veth_br_pub_indices[0], state="up")
                    else:
                        missing = []
                        if not veth_br_pub_indices:
                            missing.append(veth_br_pub)
                        if not pub_br_indices_for_master:
                            missing.append(PUB_BR)
                        raise Exception(
                            f"Required interface(s) not found for connecting {veth_br_pub} to {PUB_BR}: {', '.join(missing)}"
                        )
                except Exception as e:
                    print(
                        f"[ERROR] Failed to connect {veth_br_pub} to bridge {PUB_BR}: {e}",
                        file=sys.stderr,
                    )
                    traceback.print_exc()
                    raise

                try:
                    print(
                        f"     Configuring public IP for {veth_ns_pub} in {ns_name}..."
                    )
                    # AsyncIPRoute に変更, async with に変更
                    async with AsyncIPRoute(netns=ns_name) as ipr_ns:
                         # link_lookup を await で呼び出す
                        veth_indices_in_ns_pub = await ipr_ns.link_lookup(ifname=veth_ns_pub)
                        if veth_indices_in_ns_pub:
                            idx_veth_ns_pub = veth_indices_in_ns_pub[0]
                             # addr を await で呼び出す
                            await ipr_ns.addr(
                                "add",
                                index=idx_veth_ns_pub,
                                address=pub_ip_str,
                                mask=int(SUBNET_MASK[1:]),
                            )
                             # link を await で呼び出す
                            await ipr_ns.link("set", index=idx_veth_ns_pub, state="up")
                            print(
                                f"     Configured public {veth_ns_pub} in {ns_name} with IP {pub_ip_full}"
                            )
                        else:
                            print(
                                f"Warning: Interface {veth_ns_pub} not found in namespace {ns_name} for IP config."
                            )
                except Exception as e:
                    print(
                        f"[ERROR] Failed to configure public IP for {veth_ns_pub} in namespace {ns_name}: {e}",
                        file=sys.stderr,
                    )
                    traceback.print_exc()
                    raise
                print(
                    f"   Namespace {ns_name} configured. Private IP: {priv_ip_full}, Public IP: {pub_ip_full}."
                )

            # 3. パブリック用のネットワーク名前空間を作成 (ns-pub)
            print(
                "3. Creating public namespace (ns-pub) and connecting it to the public bridge..."
            )
            pub_ns_name = PUB_NS
            pub_ns_ip_str = f"{PUB_NET}.100"
            pub_ns_ip_full = f"{pub_ns_ip_str}{SUBNET_MASK}"

            print(f"   - Creating namespace {pub_ns_name}...")
            # create_namespace は同期のまま
            create_namespace(pub_ns_name)
            print(f"     Waiting for namespace {pub_ns_name} to be ready...")

            print(f"     Creating veth pair {VETH_PUB_NS} <--> {VETH_PUB_BR}...")
             # link_lookup を await で呼び出す
            if not await ipr.link_lookup(ifname=VETH_PUB_NS) and not await ipr.link_lookup(
                ifname=VETH_PUB_BR
            ):
                try:
                    # link を await で呼び出す
                    await ipr.link("add", ifname=VETH_PUB_NS, peer=VETH_PUB_BR, kind="veth")
                     # link_lookup を await で呼び出す
                    veth_pub_ns_idx = (await ipr.link_lookup(ifname=VETH_PUB_NS))[0]
                     # link_lookup を await で呼び出す
                    veth_pub_br_idx = (await ipr.link_lookup(ifname=VETH_PUB_BR))[0]
                    # link を await で呼び出す
                    await ipr.link("set", index=veth_pub_ns_idx, state="up")
                     # link を await で呼び出す
                    await ipr.link("set", index=veth_pub_br_idx, state="up")
                except Exception as e:
                    print(
                        f"[ERROR] Failed to create veth pair {VETH_PUB_NS}<->{VETH_PUB_BR}: {e}",
                        file=sys.stderr,
                    )
                    traceback.print_exc()
                    raise

            try:
                print(f"     Moving {VETH_PUB_NS} to namespace {pub_ns_name}...")
                 # link_lookup を await で呼び出す
                veth_pub_ns_indices_root = await ipr.link_lookup(ifname=VETH_PUB_NS)
                if veth_pub_ns_indices_root:
                     # link を await で呼び出す
                    await ipr.link(
                        "set", index=veth_pub_ns_indices_root[0], net_ns_fd=pub_ns_name
                    )
                    print(
                        f"     Successfully moved {VETH_PUB_NS} to namespace {pub_ns_name}"
                    )
                else:
                    # AsyncIPRoute に変更, async with に変更
                    async with AsyncIPRoute(netns=pub_ns_name) as ipr_check_ns:
                         # link_lookup を await で呼び出す
                        if not await ipr_check_ns.link_lookup(ifname=VETH_PUB_NS):
                            raise Exception(
                                f"{VETH_PUB_NS} not found in root or target namespace {pub_ns_name}"
                            )
                        print(f"     {VETH_PUB_NS} already in namespace {pub_ns_name}.")
            except Exception as e:
                print(
                    f"[ERROR] Failed to move or verify {VETH_PUB_NS} in namespace {pub_ns_name}: {e}",
                    file=sys.stderr,
                )
                traceback.print_exc()
                raise

            try:
                # link_lookup を await で呼び出す
                veth_pub_br_indices = await ipr.link_lookup(ifname=VETH_PUB_BR)
                 # link_lookup を await で呼び出す
                pub_br_indices_for_master_pub = await ipr.link_lookup(ifname=PUB_BR)
                if veth_pub_br_indices and pub_br_indices_for_master_pub:
                    # link を await で呼び出す
                    await ipr.link(
                        "set",
                        index=veth_pub_br_indices[0],
                        master=pub_br_indices_for_master_pub[0],
                    )
                     # link を await で呼び出す
                    await ipr.link("set", index=veth_pub_br_indices[0], state="up")
                else:
                    missing = []
                    if not veth_pub_br_indices:
                        missing.append(VETH_PUB_BR)
                    if not pub_br_indices_for_master_pub:
                        missing.append(PUB_BR)
                    raise Exception(
                        f"Required interface(s) not found for connecting {VETH_PUB_BR} to {PUB_BR}: {', '.join(missing)}"
                    )
            except Exception as e:
                print(
                    f"[ERROR] Failed to connect {VETH_PUB_BR} to bridge {PUB_BR}: {e}",
                    file=sys.stderr,
                )
                traceback.print_exc()
                raise

            print(f"     Configuring IP for {VETH_PUB_NS} in {pub_ns_name}...")
            try:
                 # AsyncIPRoute に変更, async with に変更
                async with AsyncIPRoute(netns=pub_ns_name) as ipr_ns_pub:
                    # link_lookup を await で呼び出す
                    veth_indices_in_ns_pub_final = await ipr_ns_pub.link_lookup(
                        ifname=VETH_PUB_NS
                    )
                    if veth_indices_in_ns_pub_final:
                        idx_veth_ns_pub_final = veth_indices_in_ns_pub_final[0]
                         # addr を await で呼び出す
                        await ipr_ns_pub.addr(
                            "add",
                            index=idx_veth_ns_pub_final,
                            address=pub_ns_ip_str,
                            mask=int(SUBNET_MASK[1:]),
                        )
                         # link を await で呼び出す
                        await ipr_ns_pub.link("set", index=idx_veth_ns_pub_final, state="up")
                        print(
                            f"     Configured {VETH_PUB_NS} with IP {pub_ns_ip_full} in {pub_ns_name}"
                        )
                    else:
                        print(
                            f"Warning: Interface {VETH_PUB_NS} not found in namespace {pub_ns_name} for IP config."
                        )
            except Exception as e:
                print(
                    f"[ERROR] Failed to configure public IP for {VETH_PUB_NS} in namespace {pub_ns_name}: {e}",
                    file=sys.stderr,
                )
                traceback.print_exc()
                raise
            print(
                f"   Public namespace {pub_ns_name} configured with IP {pub_ns_ip_full}."
            )

        print("\nSetup complete.")
        print(
            f"Server namespaces: {', '.join([f'ns{i}' for i in range(1, NUM_SERVERS + 1)])}"
        )
        print(f"Public namespace : {PUB_NS}")
        print(f"Private network  : {PRIV_NET}.0/24 via bridge {PRIV_BR}")
        print(f"Public network   : {PUB_NET}.0/24 via bridge {PUB_BR}")

    except NetlinkError as e:
        print(f"\nError during network setup: {e}", file=sys.stderr)
        print("Stack trace:", file=sys.stderr)
        traceback.print_exc()
        print("Attempting cleanup...", file=sys.stderr)
        # cleanup_network も非同期にする必要あり
        await cleanup_network()
        sys.exit(1)
    except Exception as e:
        print(f"\nAn unexpected error occurred: {e}", file=sys.stderr)
        print("Stack trace:", file=sys.stderr)
        traceback.print_exc()
        print("Attempting cleanup...", file=sys.stderr)
        # cleanup_network も非同期にする必要あり
        await cleanup_network()
        sys.exit(1)


async def cleanup_network():
    """
    作成したネットワーク構成を削除する (非同期版)
    """
    print("\nCleaning up network simulation...")

    try:
        namespaces_to_delete = [f"ns{i}" for i in range(1, NUM_SERVERS + 1)] + [PUB_NS]
        for ns_name in namespaces_to_delete:
            print(f"   - Deleting namespace {ns_name}...")
             # remove_namespace は同期のまま
            if remove_namespace(ns_name):
                print(f"     Namespace {ns_name} deleted.")

        async with AsyncIPRoute() as ipr: # async with に変更
            print("   - Deleting veth interfaces (peers in root namespace)...")
            veth_peers_in_root_ns = []
            for i in range(1, NUM_SERVERS + 1):
                veth_peers_in_root_ns.append(f"veth-br{i}-priv")
                veth_peers_in_root_ns.append(f"veth-br{i}-pub")
            veth_peers_in_root_ns.append(VETH_PUB_BR)

            for veth_ifname in veth_peers_in_root_ns:
                try:
                     # link_lookup を await で呼び出す
                    indices = await ipr.link_lookup(ifname=veth_ifname)
                    if indices:
                        print(f"     - Deleting veth interface {veth_ifname}...")
                        # link を await で呼び出す
                        await ipr.link("del", index=indices[0])
                except NetlinkError as e:
                    if e.code == errno.ENOENT:
                        print(
                            f"       Veth {veth_ifname} already deleted or not found (ENOENT)."
                        )
                    else:
                        print(
                            f"       Warning: Could not delete veth {veth_ifname}: {e}",
                            file=sys.stderr,
                        )
                except Exception as e:
                    print(
                        f"       Warning: Could not delete veth {veth_ifname}: {e}",
                        file=sys.stderr,
                    )

            print(
                "   - Searching for and deleting any other remaining veth interfaces in root ns..."
            )
            try:
                # get_links を await で呼び出す
                async for link in await ipr.get_links():
                    ifname = link.get_attr("IFLA_IFNAME")
                    kind = None
                    link_info = link.get_attr("IFLA_LINKINFO")
                    if link_info:
                        info_data = link_info.get_attr("IFLA_INFO_DATA")
                        if info_data:
                            kind = info_data.get_attr("IFLA_INFO_KIND")

                    if ifname and kind == "veth":
                        try:
                            print(
                                f"     - Deleting remaining veth interface {ifname} (index {link['index']})..."
                            )
                            # link を await で呼び出す
                            await ipr.link("del", ifname=ifname)
                        except NetlinkError as e_del:
                            if e_del.code == errno.ENOENT:
                                print(
                                    f"       Veth {ifname} disappeared before deletion (ENOENT)."
                                )
                            else:
                                print(
                                    f"       Warning: Could not delete veth {ifname}: {e_del}",
                                    file=sys.stderr,
                                )
                        except Exception as e_del_other:
                            print(
                                f"       Warning: Could not delete veth {ifname}: {e_del_other}",
                                file=sys.stderr,
                            )
            except Exception as e_get_links:
                print(
                    f"[ERROR] Error while listing links for remaining veth cleanup: {e_get_links}",
                    file=sys.stderr,
                )

            bridges_to_delete = [PRIV_BR, PUB_BR]
            for br_name in bridges_to_delete:
                try:
                    # link_lookup を await で呼び出す
                    br_indices = await ipr.link_lookup(ifname=br_name)
                    if br_indices:
                        print(f"   - Setting bridge {br_name} down...")
                         # link を await で呼び出す
                        await ipr.link("set", index=br_indices[0], state="down")
                        print(f"   - Deleting bridge {br_name}...")
                         # link を await で呼び出す
                        await ipr.link("del", index=br_indices[0])
                        print(f"     Bridge {br_name} deleted.")
                except NetlinkError as e_br:
                    if e_br.code == errno.ENOENT:
                        print(
                            f"     Bridge {br_name} not found (ENOENT) during cleanup steps."
                        )
                    else:
                        print(
                            f"     Warning: Error during cleanup of bridge {br_name}: {e_br}",
                            file=sys.stderr,
                        )
                except Exception as e_br_other:
                    print(
                        f"     Warning: Error during cleanup of bridge {br_name}: {e_br_other}",
                        file=sys.stderr,
                    )

        print("Cleanup complete.")

    except NetlinkError as e:
        print(f"\nError during cleanup: {e}", file=sys.stderr)
    except Exception as e:
        print(f"\nAn unexpected error occurred during cleanup: {e}", file=sys.stderr)


async def add_delay_to_interface(ns_name, if_name, delay_ms):
    """
    指定された名前空間のインターフェースに netem 遅延を追加する (非同期版)
    Equivalent to: ip netns exec <ns_name> tc qdisc add dev <if_name> root netem delay <delay_ms>ms
    """
    print(f"\nAdding {delay_ms}ms delay to {if_name} in namespace {ns_name}...")
    try:
         # AsyncIPRoute に変更, async with に変更
        async with AsyncIPRoute(netns=ns_name) as ipr_ns:
            # link_lookup を await で呼び出す
            indices = await ipr_ns.link_lookup(ifname=if_name)
            if not indices:
                print(f"[ERROR] Interface {if_name} not found in namespace {ns_name}", file=sys.stderr)
                return
            if_index = indices[0]

            # 遅延をミリ秒からマイクロ秒に変換
            delay_us = delay_ms * 1000

            # tc qdisc add dev <if_index> handle ffff: root netem delay <delay_us>
            # handle 0xffff0000 は root qdisc を意味します。
            # pyroute2 v0.7.0 以降では、キーワード引数でオプションを渡すことが推奨されます。
             # tc を await で呼び出す
            await ipr_ns.tc(
                command='add',
                kind='netem',
                index=if_index,
                handle=TC_H_ROOT, # ROOT handle
                # netem 固有のオプションはキーワード引数として渡す
                delay=delay_us
                # 他の netem オプション (例: jitter, loss) もここに追加可能
                # jitter=5000, # 5ms jitter
                # loss=10.0,  # 10% packet loss
            )
            print(f"   Successfully added {delay_ms}ms delay to {if_name} (index {if_index}) in {ns_name}.")

    except NetlinkError as e:
        # 例えば、qdisc がすでに存在する場合 (errno.EEXIST)
        if e.code == errno.EEXIST:
             print(f"   Warning: qdisc already exists on {if_name} in {ns_name}. Trying 'replace'.", file=sys.stderr)
             try:
                 # 既存の qdisc を置き換える試み
                 # AsyncIPRoute に変更, async with に変更
                 async with AsyncIPRoute(netns=ns_name) as ipr_replace:
                     # tc を await で呼び出す
                     await ipr_replace.tc(
                         command='replace', # 'add' の代わりに 'replace'
                         kind='netem',
                         index=if_index,
                         handle=TC_H_ROOT,
                         delay=delay_us
                     )
                     print(f"   Successfully replaced existing qdisc on {if_name} in {ns_name} with {delay_ms}ms delay.")
             except Exception as e_replace:
                 print(f"[ERROR] Failed to replace qdisc on {if_name} in {ns_name}: {e_replace}", file=sys.stderr)
                 traceback.print_exc()
        else:
            print(f"[ERROR] Failed to add/replace qdisc on {if_name} in {ns_name}: {e}", file=sys.stderr)
            traceback.print_exc()
    except Exception as e:
        print(f"[ERROR] An unexpected error occurred while adding delay to {if_name} in {ns_name}: {e}", file=sys.stderr)
        traceback.print_exc()

async def delete_delay_from_interface(ns_name, if_name):
    """
    指定された名前空間のインターフェースから netem 遅延を削除する (非同期版)
    Equivalent to: ip netns exec <ns_name> tc qdisc del dev <if_name> root netem
    """
    print(f"\nRemoving delay from {if_name} in namespace {ns_name}...")

    try:
        # AsyncIPRoute に変更, async with に変更
        async with AsyncIPRoute(netns=ns_name) as ipr_ns:
             # link_lookup を await で呼び出す
            indices = await ipr_ns.link_lookup(ifname=if_name)
            if not indices:
                print(f"[ERROR] Interface {if_name} not found in namespace {ns_name}", file=sys.stderr)
                return
            if_index = indices[0]

            # 既存の qdisc を取得
            # get_qdiscs を await で呼び出す
            qdiscs = await ipr_ns.get_qdiscs(index=if_index)
            netem_exists = False
            for q in qdiscs:
                # root ハンドル (0xffff0000) かつ kind が netem かチェック
                if q.get('handle') == 0xffff0000 and q.get_attr('TCA_KIND') == 'netem':
                    netem_exists = True
                    break

            if netem_exists:
                # netem qdisc が存在する場合のみ削除を実行
                print(f"   Found existing netem qdisc on {if_name}. Attempting removal...")
                 # tc を await で呼び出す
                await ipr_ns.tc(
                    command='del',
                    kind='netem',
                    index=if_index,
                    handle=0xffff0000
                )
                print(f"   Successfully removed delay from {if_name} in {ns_name}.")
            else:
                # netem qdisc が存在しない場合は何もしない
                print(f"   No root netem qdisc found on {if_name} in {ns_name}. No removal needed.")


    except NetlinkError as e:
        # 削除中に予期せぬ Netlink エラーが発生した場合
        print(f"[ERROR] Failed to delete delay from {if_name} in {ns_name}: {e}", file=sys.stderr)

# --- メイン処理 ---

# main ブロックを async 関数でラップし、asyncio.run で実行
async def main():
    if os.geteuid() != 0:
        print("This script requires root privileges. Run with sudo.", file=sys.stderr)
        sys.exit(1)

    if len(sys.argv) != 2 or sys.argv[1] not in ["create", "cleanup", "add_delay", "delete_delay"]:
        print("Usage: sudo python3 main.py [create|cleanup|add_delay|delete_delay]", file=sys.stderr)
        sys.exit(1)

    action = sys.argv[1]

    if action == "create":
        await create_network()

    elif action == "add_delay":
        target_namespace = "ns2"
        target_interface = "veth-ns2-pub"
        delay_milliseconds = 100
        await add_delay_to_interface(target_namespace, target_interface, delay_milliseconds)

    elif action == "delete_delay":
        target_namespace = "ns2"
        target_interface = "veth-ns2-pub"
        await delete_delay_from_interface(target_namespace, target_interface)

    elif action == "cleanup":
        await cleanup_network()

    if action == "create":
        print("\n--- Verification Commands ---")
        print("  # List all network namespaces")
        print("  sudo ip netns list")
        print("\n  # Show details of the private bridge")
        print(f"  sudo ip a show dev {PRIV_BR}")
        print("\n  # Show details of the public bridge")
        print(f"  sudo ip a show dev {PUB_BR}")
        print(
            "\n  # Show IP addresses in server namespace ns1 (replace ns1 with ns2, ns3 etc. as needed)"
        )
        print("  sudo ip netns exec ns1 ip a")
        print("\n  # Show IP addresses in the public namespace")
        print(f"  sudo ip netns exec {PUB_NS} ip a")
        print(
            f"\n  # Ping from ns1 to ns2 on the private network (assumes ns2 is {PRIV_NET}.2)"
        )
        print(f"  sudo ip netns exec ns1 ping -c 3 {PRIV_NET}.2")
        print(
            f"\n  # Ping from the public namespace to ns1 on the public network (assumes ns1 public IP is {PUB_NET}.1)"
        )
        print(f"  sudo ip netns exec {PUB_NS} ping -c 3 {PUB_NET}.1")

if __name__ == "__main__":
    asyncio.run(main()) # asyncio.run で main() を実行