import <nixpkgs/nixos/tests/make-test-python.nix> (
  { pkgs, lib }:
  let
    makeCommonConfig = hostname: {
      virtualisation.restrictNetwork = true;
      networking.useDHCP = false;
      networking.firewall.enable = false;
      networking.nftables.enable = false;
      environment.systemPackages = [ pkgs.traceroute ];
    };

    # VLANS:
    # 1 -- simulates the internal network
    # 2 -- simulates the external network
    routerBase =
      nodes:
      lib.mkMerge [
        (makeCommonConfig "router")
        {
          virtualisation.vlans = [
            1
            2
          ];
          boot.kernel.sysctl = {
            "net.ipv4.conf.all.forwarding" = lib.mkOverride 99 true;
          };
        }
      ];
  in
  {
    name = "nat-customtest";

    nodes = {
      client =
        { pkgs, nodes, ... }:
        lib.mkMerge [
          (makeCommonConfig "client")
          {
            virtualisation.vlans = [ 1 ];
            networking.defaultGateway =
              (pkgs.lib.head nodes.router.config.networking.interfaces.eth1.ipv4.addresses).address;
          }
        ];

      router =
        { nodes, ... }:
        lib.mkMerge [
          (routerBase nodes)
        ];

      server =
        { nodes, ... }:
        lib.mkMerge [
          (makeCommonConfig "server")
          {
            virtualisation.vlans = [ 2 ];
            networking.defaultGateway =
              (pkgs.lib.head nodes.router.config.networking.interfaces.eth2.ipv4.addresses).address;
          }
        ];
    };

    testScript =
      { nodes, ... }:
      let
        routerDummyNoNatClosure = nodes.routerDummyNoNat.config.system.build.toplevel;
        routerClosure = nodes.router.config.system.build.toplevel;
        clientIp = (pkgs.lib.head nodes.client.config.networking.interfaces.eth1.ipv4.addresses).address;
        serverIp = (pkgs.lib.head nodes.server.config.networking.interfaces.eth1.ipv4.addresses).address;
        routerIp = (pkgs.lib.head nodes.router.config.networking.interfaces.eth2.ipv4.addresses).address;
      in
      ''
        def wait_for_machine(m):
          m.wait_for_unit("network.target")

        client.start()
        router.start()
        server.start()

        wait_for_machine(router)
        wait_for_machine(client)
        wait_for_machine(server)

        server.succeed("traceroute client 1>&2")
        server.succeed("ip addr 1>&2")
        client.succeed("ip addr 1>&2")
        router.succeed("ip addr 1>&2")

        # Try to make a connection from the client to the server
        client.fail("timeout 3 ping -c 1 server")
        # Try to make a connection from the server to the client
        server.fail("timeout 3 ping -c 1 client")
      '';
  }
)