firewall_containers/network-go/docker/docker.go

package docker

import (
	"context"
	"fmt"
	"net"
	"os/exec"
	"strings"
	"time"

	"github.com/docker/docker/api/types"
	"github.com/docker/docker/api/types/container"
	"github.com/docker/docker/api/types/filters"
	"github.com/docker/docker/api/types/network"
	"github.com/docker/docker/client"

	"firewall_containers/network-go/config"
)

// DockerAPI defines the interface for Docker operations, enabling mock implementations for testing
type DockerAPI interface {
	Close() error
	EnsureNetwork(ctx context.Context, netCfg config.NetworkConfig) error
	RemoveNetwork(ctx context.Context, networkName string) error
	ConnectContainer(ctx context.Context, containerName, networkName, ip string) error
	DisconnectContainer(ctx context.Context, containerName, networkName string) error
	InspectContainer(ctx context.Context, containerName string) (*types.ContainerJSON, error)
	WaitForContainerRunning(ctx context.Context, containerName string, timeout time.Duration) error
	GetContainerPID(ctx context.Context, containerName string) (int, error)
	AddRouteInContainer(ctx context.Context, containerName, network, gateway string) error
	FindContainerName(ctx context.Context, name, selector string) (string, error)
}

// Client wraps the Docker SDK client
type Client struct {
	cli *client.Client
}

// Ensure Client implements DockerAPI
var _ DockerAPI = (*Client)(nil)

// NewClient creates a new Docker client
func NewClient() (*Client, error) {
	cli, err := client.NewClientWithOpts(client.FromEnv, client.WithAPIVersionNegotiation())
	if err != nil {
		return nil, fmt.Errorf("failed to create Docker client: %w", err)
	}
	return &Client{cli: cli}, nil
}

// Close closes the Docker client
func (c *Client) Close() error {
	return c.cli.Close()
}

// EnsureNetwork creates a Docker network if it does not already exist
func (c *Client) EnsureNetwork(ctx context.Context, netCfg config.NetworkConfig) error {
	existingNetworks, err := c.cli.NetworkList(ctx, network.ListOptions{
		Filters: filters.NewArgs(filters.Arg("name", netCfg.NetworkName)),
	})
	if err != nil {
		return fmt.Errorf("failed to list networks: %w", err)
	}

	for _, n := range existingNetworks {
		if n.Name == netCfg.NetworkName {
			return nil
		}
	}

	_, ipNet, err := net.ParseCIDR(netCfg.Subnet)
	if err != nil {
		return fmt.Errorf("failed to parse subnet %s: %w", netCfg.Subnet, err)
	}

	gatewayIP := net.ParseIP(netCfg.Gateway)
	if gatewayIP == nil {
		return fmt.Errorf("failed to parse gateway IP %s", netCfg.Gateway)
	}

	createOpts := network.CreateOptions{
		Driver: "bridge",
		IPAM: &network.IPAM{
			Config: []network.IPAMConfig{
				{
					Subnet:  ipNet.String(),
					Gateway: gatewayIP.String(),
				},
			},
		},
		Labels: map[string]string{
			"managed-by": "firewall-network-go",
		},
		Attachable: true,
	}

	resp, err := c.cli.NetworkCreate(ctx, netCfg.NetworkName, createOpts)
	if err != nil {
		return fmt.Errorf("failed to create network %s: %w", netCfg.NetworkName, err)
	}

	_ = resp
	return nil
}

// RemoveNetwork removes a Docker network
func (c *Client) RemoveNetwork(ctx context.Context, networkName string) error {
	err := c.cli.NetworkRemove(ctx, networkName)
	if err != nil {
		return fmt.Errorf("failed to remove network %s: %w", networkName, err)
	}
	return nil
}

// ConnectContainer connects a container to a network with a specific IP
func (c *Client) ConnectContainer(ctx context.Context, containerName, networkName, ip string) error {
	endpointSettings := &network.EndpointSettings{
		IPAMConfig: &network.EndpointIPAMConfig{
			IPv4Address: ip,
		},
	}

	err := c.cli.NetworkConnect(ctx, networkName, containerName, endpointSettings)
	if err != nil {
		return fmt.Errorf("failed to connect container %s to network %s: %w", containerName, networkName, err)
	}
	return nil
}

// DisconnectContainer disconnects a container from a network
func (c *Client) DisconnectContainer(ctx context.Context, containerName, networkName string) error {
	err := c.cli.NetworkDisconnect(ctx, networkName, containerName, true)
	if err != nil {
		return fmt.Errorf("failed to disconnect container %s from network %s: %w", containerName, networkName, err)
	}
	return nil
}

// InspectContainer returns the container's details
func (c *Client) InspectContainer(ctx context.Context, containerName string) (*types.ContainerJSON, error) {
	container, err := c.cli.ContainerInspect(ctx, containerName)
	if err != nil {
		return nil, fmt.Errorf("failed to inspect container %s: %w", containerName, err)
	}
	return &container, nil
}

// WaitForContainerRunning waits until a container is in running state, with a timeout
func (c *Client) WaitForContainerRunning(ctx context.Context, containerName string, timeout time.Duration) error {
	ctx, cancel := context.WithTimeout(ctx, timeout)
	defer cancel()

	ticker := time.NewTicker(500 * time.Millisecond)
	defer ticker.Stop()

	for {
		select {
		case <-ctx.Done():
			return fmt.Errorf("timeout waiting for container %s to be running", containerName)
		case <-ticker.C:
			container, err := c.cli.ContainerInspect(ctx, containerName)
			if err != nil {
				continue
			}
			if container.State != nil && container.State.Running {
				return nil
			}
		}
	}
}

// GetContainerPID returns the PID of a container for nsenter operations
func (c *Client) GetContainerPID(ctx context.Context, containerName string) (int, error) {
	cont, err := c.cli.ContainerInspect(ctx, containerName)
	if err != nil {
		return 0, fmt.Errorf("failed to inspect container %s: %w", containerName, err)
	}
	if cont.State == nil || !cont.State.Running {
		return 0, fmt.Errorf("container %s is not running", containerName)
	}
	return cont.State.Pid, nil
}

// AddRouteInContainer adds a route inside a container's network namespace using nsenter
func (c *Client) AddRouteInContainer(ctx context.Context, containerName, network, gateway string) error {
	pid, err := c.GetContainerPID(ctx, containerName)
	if err != nil {
		return fmt.Errorf("failed to get PID for container %s: %w", containerName, err)
	}

	args := []string{
		"-t", fmt.Sprintf("%d", pid),
		"-n", "--",
		"ip", "route", "add", network, "via", gateway,
	}

	cmd := exec.Command("nsenter", args...)
	output, err := cmd.CombinedOutput()
	if err != nil {
		return fmt.Errorf("failed to add route in container %s: %w\noutput: %s", containerName, err, string(output))
	}
	return nil
}

// FindContainerName attempts to find a running container by name or selector.
// First tries the exact name, then tries listing running containers whose name
// starts with the selector prefix (or the name prefix), matching the old shell
// script's grep $D"-" behavior.
func (c *Client) FindContainerName(ctx context.Context, name, selector string) (string, error) {
	// First try the exact name
	cont, err := c.cli.ContainerInspect(ctx, name)
	if err == nil && cont.State != nil && cont.State.Running {
		return name, nil
	}

	// Try exact selector
	if selector != "" && selector != name {
		cont, err := c.cli.ContainerInspect(ctx, selector)
		if err == nil && cont.State != nil && cont.State.Running {
			return selector, nil
		}
	}

	// Try prefix matching with selector (old shell script behavior: grep $D"-")
	candidates := []string{name, selector}
	for _, candidate := range candidates {
		if candidate == "" {
			continue
		}
		// Extract prefix before first dash if present
		prefix := candidate
		if strings.Contains(candidate, "-") {
			prefix = candidate[:strings.Index(candidate, "-")]
		}

		containers, err := c.cli.ContainerList(ctx, container.ListOptions{})
		if err != nil {
			continue
		}

		for _, container := range containers {
			// Remove leading / from container names
			for _, cName := range container.Names {
				cName = strings.TrimPrefix(cName, "/")
				if strings.HasPrefix(cName, prefix+"-") && container.State == "running" {
					return cName, nil
				}
			}
		}
	}

	return "", fmt.Errorf("no running container found matching name=%q selector=%q", name, selector)
}