vpn-provisioner/src/main/java/com/litoralregas/vpnprovisioner/vps/WireGuardVpsService.java

package com.litoralregas.vpnprovisioner.vps;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.litoralregas.vpnprovisioner.vps.dto.VpsHealthResponse;
import org.springframework.stereotype.Service;

import java.util.HashSet;
import java.util.Set;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

@Service
public class WireGuardVpsService {

    private static final Pattern VPN_IP_PATTERN =
            Pattern.compile("\\b198\\.19\\.\\d{1,3}\\.\\d{1,3}\\b");

    private final SshService sshService;
    private final ObjectMapper objectMapper;

    public WireGuardVpsService(SshService sshService, ObjectMapper objectMapper) {
        this.sshService = sshService;
        this.objectMapper = objectMapper;
    }

    public Set<String> findUsedVpnIps() {
        SshCommandResult result = sshService.executeOnConfiguredVps(
                "sudo /usr/local/sbin/lr-wg-used-ips"
        );

        if (result.exitCode() != 0) {
            throw new SshCommandException(
                    "Failed to query WireGuard allowed IPs: " + result.stderr()
            );
        }

        return parseVpnIps(result.stdout());
    }

    Set<String> parseVpnIps(String output) {
        Set<String> ips = new HashSet<>();

        Matcher matcher = VPN_IP_PATTERN.matcher(output);

        while (matcher.find()) {
            ips.add(matcher.group());
        }

        return ips;
    }

    public WireGuardPeerApplyResult applyPeer(String publicKey, String allowedIps) {

        String command = """
            sudo /usr/local/sbin/lr-wg-add-peer '%s' '%s'
            """.formatted(publicKey, allowedIps);

        SshCommandResult result = sshService.executeOnConfiguredVps(command);

        if (result.exitCode() != 0) {
            throw new SshCommandException(
                    "Failed to apply WireGuard peer: " + result.stderr()
            );
        }

        return new WireGuardPeerApplyResult(
                publicKey,
                allowedIps,
                true
        );
    }

    public VpsHealthResponse getVpsHealth() {
        SshCommandResult result = sshService.executeOnConfiguredVps(
                "sudo /usr/local/sbin/lr-vps-health"
        );

        if (result.exitCode() != 0) {
            throw new SshCommandException(
                    "Failed to query VPS health: " + result.stderr()
            );
        }

        try {
            return objectMapper.readValue(result.stdout(), VpsHealthResponse.class);
        } catch (JsonProcessingException e) {
            throw new IllegalStateException(
                    "Invalid VPS health JSON returned by script",
                    e
            );
        }
    }

    public String restoreLastWireGuardBackup() {
        SshCommandResult result = sshService.executeOnConfiguredVps(
                "sudo /usr/local/sbin/lr-wg-restore-last-backup"
        );

        if (result.exitCode() != 0) {
            throw new SshCommandException(
                    "Failed to restore WireGuard backup: " + result.stderr()
            );
        }

        return result.stdout();
    }

    public String showAllowedIps() {
        SshCommandResult result = sshService.executeOnConfiguredVps(
                "sudo /usr/local/sbin/lr-wg-used-ips"
        );

        if (result.exitCode() != 0) {
            throw new SshCommandException(
                    "Failed to query WireGuard allowed IPs: " + result.stderr()
            );
        }

        return result.stdout();
    }
}