usky-modules/service-ems/service-ems-biz/src/main/java/com/usky/ems/protocol/TcpClient.java

package com.usky.ems.protocol;

import com.usky.ems.util.AesUtil;
import org.dom4j.Document;
import org.dom4j.DocumentHelper;
import org.dom4j.Element;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.net.Socket;
import java.util.concurrent.atomic.AtomicBoolean;

/**
 * TCP客户端
 * 用于与能耗监管系统建立连接并发送数据
 *
 * @author system
 * @since 2024-01-01
 */
public class TcpClient {

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

    private String host;
    private int port;
    private String authKey;
    private Socket socket;
    private InputStream inputStream;
    private OutputStream outputStream;
    private AtomicBoolean connected = new AtomicBoolean(false);
    private AtomicBoolean authenticated = new AtomicBoolean(false);

    public TcpClient(String host, int port, String authKey) {
        this.host = host;
        this.port = port;
        this.authKey = authKey;
    }

    /**
     * 建立TCP连接
     *
     * @return 是否连接成功
     */
    public boolean connect() {
        try {
            if (socket != null && !socket.isClosed()) {
                socket.close();
            }

            socket = new Socket(host, port);
            socket.setSoTimeout(30000); // 30秒超时
            inputStream = socket.getInputStream();
            outputStream = socket.getOutputStream();
            connected.set(true);
            authenticated.set(false);

            logger.info("TCP连接成功: {}:{}", host, port);
            return true;
        } catch (IOException e) {
            logger.error("TCP连接失败: {}:{}", host, port, e);
            connected.set(false);
            return false;
        }
    }

    /**
     * 关闭连接
     */
    public void close() {
        try {
            if (inputStream != null) {
                inputStream.close();
            }
            if (outputStream != null) {
                outputStream.close();
            }
            if (socket != null && !socket.isClosed()) {
                socket.close();
            }
            connected.set(false);
            authenticated.set(false);
            logger.info("TCP连接已关闭");
        } catch (IOException e) {
            logger.error("关闭TCP连接失败", e);
        }
    }

    /**
     * 身份认证
     *
     * @param buildingId 建筑ID
     * @param gatewayId  网关ID
     * @return 是否认证成功
     */
    public boolean authenticate(String buildingId, String gatewayId) {
        if (!connected.get()) {
            logger.error("未建立TCP连接，无法进行身份认证");
            return false;
        }

        try {
            // 1. 发送身份认证请求
            String requestXml = com.usky.ems.util.XmlBuilder.buildAuthRequest(buildingId, gatewayId);
            sendPacket(NetworkPacket.TYPE_AUTH, requestXml, false);

            // 2. 接收服务端返回的随机序列
            NetworkPacket response = receivePacket();
            if (response == null || response.getType() != NetworkPacket.TYPE_AUTH) {
                logger.error("接收身份认证响应失败");
                return false;
            }

            // 认证响应不加密
            String responseXml = new String(response.getData(), java.nio.charset.StandardCharsets.UTF_8);
            String sequence = parseSequenceFromXml(responseXml);
            if (sequence == null || sequence.isEmpty()) {
                logger.error("解析随机序列失败");
                return false;
            }

            // 3. 计算MD5：密钥 + 随机序列
            String md5 = com.usky.ems.util.Md5Util.md5(authKey + sequence);

            // 4. 发送MD5值
            String md5Xml = com.usky.ems.util.XmlBuilder.buildAuthMd5Request(buildingId, gatewayId, md5);
            sendPacket(NetworkPacket.TYPE_AUTH, md5Xml, false);

            // 5. 接收认证结果
            NetworkPacket result = receivePacket();
            if (result == null || result.getType() != NetworkPacket.TYPE_AUTH) {
                logger.error("接收认证结果失败");
                return false;
            }

            // 认证响应不加密
            String resultXml = new String(result.getData(), java.nio.charset.StandardCharsets.UTF_8);
            boolean success = parseAuthResult(resultXml);

            if (success) {
                authenticated.set(true);
                logger.info("身份认证成功");
            } else {
                logger.error("身份认证失败");
            }

            return success;
        } catch (Exception e) {
            logger.error("身份认证过程异常", e);
            return false;
        }
    }

    /**
     * 发送心跳
     *
     * @param buildingId 建筑ID
     * @param gatewayId  网关ID
     * @return 是否发送成功
     */
    public boolean sendHeartbeat(String buildingId, String gatewayId) {
        if (!connected.get() || !authenticated.get()) {
            logger.warn("连接未建立或未认证，无法发送心跳");
            return false;
        }

        try {
            String heartbeatXml = com.usky.ems.util.XmlBuilder.buildHeartbeatRequest(buildingId, gatewayId);
            sendPacket(NetworkPacket.TYPE_HEARTBEAT, heartbeatXml, false);
            logger.debug("心跳发送成功");
            return true;
        } catch (Exception e) {
            logger.error("发送心跳失败", e);
            return false;
        }
    }

    /**
     * 发送能耗数据
     *
     * @param xmlData XML数据（未加密）
     * @return 是否发送成功
     */
    public boolean sendEnergyData(String xmlData) {
        if (!connected.get() || !authenticated.get()) {
            logger.warn("连接未建立或未认证，无法发送能耗数据");
            return false;
        }

        try {
            sendPacket(NetworkPacket.TYPE_ENERGY_DATA, xmlData, true);
            logger.info("能耗数据发送成功");

            // 接收服务端响应
            NetworkPacket response = receivePacket();
            if (response != null && response.getType() == NetworkPacket.TYPE_ENERGY_DATA) {
                // 能耗数据响应需要解密
                try {
                    byte[] responseData = response.getData();
                    String responseXml = null;
                    Exception lastException = null;

                    // 方法1：首先尝试直接解密字节数组
                    try {
                        responseXml = AesUtil.decrypt(authKey,responseData);
                        logger.debug("直接解密字节数组成功");
                    } catch (Exception e1) {
                        lastException = e1;
                        logger.debug("直接解密字节数组失败: {}", e1.getMessage());

                        // 方法2：如果直接解密失败，尝试作为Base64字符串处理
                        // 先检查数据是否可能是Base64字符串（只包含Base64字符）
                        if (responseXml == null) {
                            try {
                                String dataString = new String(responseData, java.nio.charset.StandardCharsets.UTF_8);
                                // 检查是否只包含Base64字符（A-Z, a-z, 0-9, +, /, =）
                                if (dataString.matches("^[A-Za-z0-9+/=]+$")) {
                                    // Base64解码后再解密
                                    byte[] decodedBytes = java.util.Base64.getDecoder().decode(dataString);
                                    responseXml = AesUtil.decrypt(authKey,decodedBytes);
                                    logger.debug("Base64解码后解密成功");
                                } else {
                                    logger.debug("数据不是Base64格式，跳过Base64解码");
                                }
                            } catch (Exception e2) {
                                lastException = e2;
                                logger.debug("Base64解码后解密失败: {}", e2.getMessage());
                            }
                        }
                    }

                    if (responseXml != null) {
                        logger.debug("服务端响应: {}", responseXml);
                    } else {
                        throw lastException != null ? lastException : new Exception("所有解密方法都失败");
                    }
                } catch (Exception e) {
                    logger.warn("解密服务端响应失败", e);
                    logger.debug("响应数据长度: {}, 前100字节: {}",
                            response.getData() != null ? response.getData().length : 0,
                            response.getData() != null && response.getData().length > 0
                                    ? java.util.Arrays.toString(java.util.Arrays.copyOf(response.getData(), Math.min(100, response.getData().length)))
                                    : "null");
                }
                return true;
            }

            return true;
        } catch (Exception e) {
            logger.error("发送能耗数据失败", e);
            return false;
        }
    }

    /**
     * 发送数据包
     *
     * @param type     消息类型
     * @param xmlData  XML数据
     * @param encrypt  是否加密
     */
    private void sendPacket(byte type, String xmlData, boolean encrypt) throws IOException {
        byte[] data;
        if (encrypt) {
            // 能耗数据需要AES加密
            try {
                data = AesUtil.encryptEnergyXml(xmlData, authKey);
            } catch (Exception e) {
                throw new IOException("加密能耗数据失败", e);
            }
        } else {
            // 身份认证和心跳不加密
            data = xmlData.getBytes(java.nio.charset.StandardCharsets.UTF_8);
        }

        NetworkPacket packet = new NetworkPacket();
        packet.setType(type);
        packet.setData(data);

        byte[] packetBytes = packet.encode();
        outputStream.write(packetBytes);
        outputStream.flush();
    }

    /**
     * 接收数据包
     *
     * @return 接收到的数据包
     */
    private NetworkPacket receivePacket() throws IOException {
        // 先读取7字节（Head + Type + Length）
        byte[] header = new byte[7];
        int bytesRead = 0;
        while (bytesRead < 7) {
            int n = inputStream.read(header, bytesRead, 7 - bytesRead);
            if (n == -1) {
                throw new IOException("连接已关闭");
            }
            bytesRead += n;
        }

        // 解析Length
        int length = ((header[3] & 0xFF) << 24) |
                ((header[4] & 0xFF) << 16) |
                ((header[5] & 0xFF) << 8) |
                (header[6] & 0xFF);

        // 读取Data
        byte[] data = new byte[length];
        bytesRead = 0;
        while (bytesRead < length) {
            int n = inputStream.read(data, bytesRead, length - bytesRead);
            if (n == -1) {
                throw new IOException("连接已关闭");
            }
            bytesRead += n;
        }

        // 组装完整数据包
        byte[] packetBytes = new byte[7 + length];
        System.arraycopy(header, 0, packetBytes, 0, 7);
        System.arraycopy(data, 0, packetBytes, 7, length);

        return NetworkPacket.decode(packetBytes);
    }

    /**
     * 从XML中解析随机序列
     */
    private String parseSequenceFromXml(String xml) {
        try {
            Document document = DocumentHelper.parseText(xml);
            Element root = document.getRootElement();
            Element idValidate = root.element("id_validate");
            if (idValidate != null) {
                Element sequence = idValidate.element("sequence");
                if (sequence != null) {
                    return sequence.getTextTrim();
                }
            }
        } catch (Exception e) {
            logger.error("解析随机序列失败", e);
        }
        return null;
    }

    /**
     * 从XML中解析认证结果
     */
    private boolean parseAuthResult(String xml) {
        try {
            Document document = DocumentHelper.parseText(xml);
            Element root = document.getRootElement();
            Element idValidate = root.element("id_validate");
            if (idValidate != null && "result".equals(idValidate.attributeValue("operation"))) {
                Element result = idValidate.element("result");
                if (result != null) {
                    return "pass".equals(result.getTextTrim());
                }
            }
        } catch (Exception e) {
            logger.error("解析认证结果失败", e);
        }
        return false;
    }

    public boolean isConnected() {
        return connected.get() && socket != null && !socket.isClosed();
    }

    public boolean isAuthenticated() {
        return authenticated.get();
    }
}