ai-security-xdr/haobang-security-xdr/syslog-serve/src/main/java/com/haobang/interlocking/ProbeHeartbeatClient.java

package com.haobang.interlocking;

import com.fasterxml.jackson.databind.ObjectMapper;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.http.*;
import org.springframework.stereotype.Component;
import org.springframework.web.client.RestTemplate;

import java.net.InetAddress;
import java.net.UnknownHostException;
import java.time.LocalDateTime;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * 探针心跳客户端
 * 定期向平台端发送心跳，支持重试和指数退避
 */
@Component
public class ProbeHeartbeatClient {

    private static final Logger logger = LoggerFactory.getLogger(ProbeHeartbeatClient.class);

    private final RestTemplate restTemplate = new RestTemplate();
    private final ObjectMapper objectMapper = new ObjectMapper();

    @Value("${probe.heartbeat.enabled:true}")
    private boolean heartbeatEnabled;

    @Value("${probe.heartbeat.interval-seconds:60}")
    private int heartbeatIntervalSeconds;

    @Value("${app.service.device_collect_id:1}")
    private String collectId;

    @Value("${app.service.device_collect_name:采集探针-01}")
    private String collectName;

    @Value("${app.service.version:V1.0.0-20260509}")
    private String appVersion;

    @Value("${probe.platform.api-url:http://localhost:8089/xdrservice/interlocking/probe/heartbeat}")
    private String platformApiUrl;

    @Value("${interlocking.api-key:a1b2c3d4e5f6g7h8i9j0k1l2m3n4o5p6}")
    private String apiKey;

    /** 重试次数 */
    private final AtomicInteger retryCount = new AtomicInteger(0);
    
    /** 最大重试次数 */
    private static final int MAX_RETRIES = 5;
    
    /** 初始重试延迟（毫秒） */
    private static final long INITIAL_BACKOFF_MS = 1000;
    
    /** 最大退避延迟（毫秒） */
    private static final long MAX_BACKOFF_MS = 60000;

    /**
     * 发送心跳
     * @return 是否成功
     */
    public boolean sendHeartbeat() {
        if (!heartbeatEnabled) {
            logger.debug("心跳发送已禁用");
            return false;
        }

        // 构建心跳数据
        Map<String, Object> heartbeatData = buildHeartbeatData();
        
        try {
            // 发送HTTP POST请求
            HttpHeaders headers = new HttpHeaders();
            headers.setContentType(MediaType.APPLICATION_JSON);
            headers.set("X-API-KEY", apiKey);
            
            HttpEntity<Map<String, Object>> entity = new HttpEntity<>(heartbeatData, headers);
            ResponseEntity<String> response = restTemplate.exchange(
                    platformApiUrl, 
                    HttpMethod.POST, 
                    entity, 
                    String.class
            );
            
            if (response.getStatusCode() == HttpStatus.OK) {
                // 重置重试计数
                retryCount.set(0);
                
                // 检查响应
                String body = response.getBody();
                if (body != null && body.contains("\"code\":200")) {
                    logger.debug("心跳发送成功: collectId={}, timestamp={}", collectId, heartbeatData.get("timestamp"));
                    return true;
                }
            }
            
            // 请求失败，准备重试
            handleFailure("HTTP " + response.getStatusCodeValue());
            return false;
            
        } catch (Exception e) {
            handleFailure(e.getMessage());
            return false;
        }
    }

    /**
     * 发送心跳（简化版，供其他服务调用）
     */
    public boolean sendHeartbeat(String alertType, String alertContent) {
        // 直接发送普通心跳
        return sendHeartbeat();
    }

    /**
     * 构建心跳数据
     */
    private Map<String, Object> buildHeartbeatData() {
        Map<String, Object> data = new HashMap<>();
        
        // 使用配置的collectId或自动获取
        String probeId = collectId;
        if (probeId == null || probeId.isEmpty()) {
            probeId = getLocalHostIdentifier();
        }
        
        data.put("collectId", probeId);
        data.put("collectName", collectName != null && !collectName.isEmpty() ? collectName : "SyslogServe");
        data.put("deviceIp", getLocalIp());
        data.put("appVersion", appVersion);
        data.put("timestamp", LocalDateTime.now().toString());
        
        // 添加负载状态（可选）
        Map<String, Object> loadStatus = new HashMap<>();
        loadStatus.put("memoryUsage", getMemoryUsage());
        loadStatus.put("threadCount", Thread.activeCount());
        data.put("loadStatus", toJson(loadStatus));
        
        return data;
    }

    /**
     * 获取本地主机标识符
     */
    private String getLocalHostIdentifier() {
        try {
            InetAddress ip = InetAddress.getLocalHost();
            return ip.getHostName() + "-" + ip.getHostAddress();
        } catch (UnknownHostException e) {
            return "unknown-" + System.currentTimeMillis();
        }
    }

    /**
     * 获取本地IP地址
     */
    private String getLocalIp() {
        try {
            InetAddress ip = InetAddress.getLocalHost();
            return ip.getHostAddress();
        } catch (UnknownHostException e) {
            return "127.0.0.1";
        }
    }

    /**
     * 获取内存使用情况
     */
    private double getMemoryUsage() {
        Runtime runtime = Runtime.getRuntime();
        long totalMemory = runtime.totalMemory();
        long freeMemory = runtime.freeMemory();
        long usedMemory = totalMemory - freeMemory;
        return Math.round((double) usedMemory / totalMemory * 100 * 100) / 100.0;
    }

    /**
     * 对象转JSON字符串
     */
    private String toJson(Object obj) {
        try {
            return objectMapper.writeValueAsString(obj);
        } catch (Exception e) {
            return "{}";
        }
    }

    /**
     * 处理发送失败，使用指数退避重试
     */
    private void handleFailure(String errorMsg) {
        int currentRetry = retryCount.incrementAndGet();
        
        if (currentRetry <= MAX_RETRIES) {
            // 计算退避延迟
            long backoffMs = Math.min(INITIAL_BACKOFF_MS * (1L << (currentRetry - 1)), MAX_BACKOFF_MS);
            
            logger.warn("心跳发送失败 (尝试 {}/{}): {}, {}ms后将重试", 
                    currentRetry, MAX_RETRIES, errorMsg, backoffMs);
            
            try {
                Thread.sleep(backoffMs);
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
            }
        } else {
            logger.error("心跳发送失败，已达到最大重试次数 {}: {}", MAX_RETRIES, errorMsg);
            // 重置重试计数，下次将重新开始
            retryCount.set(0);
        }
    }

    /**
     * 获取配置的心跳间隔
     */
    public int getHeartbeatIntervalSeconds() {
        return heartbeatIntervalSeconds;
    }

    /**
     * 获取探针ID
     */
    public String getCollectId() {
        return collectId != null && !collectId.isEmpty() ? collectId : getLocalHostIdentifier();
    }

    /**
     * 检查心跳功能是否启用
     */
    public boolean isEnabled() {
        return heartbeatEnabled;
    }
}