基于AES加密websocket连接传输

发送端：-  
/**
     * 生成AES密钥
     *
     * @return
     * @throws NoSuchAlgorithmException
     */
    private static SecretKey generateSecretKey() throws NoSuchAlgorithmException {

        KeyGenerator keyGenerator = KeyGenerator.getInstance("AES");

        keyGenerator.init(256); // 使用AES-256

        return keyGenerator.generateKey();

    }
    /**
     * 加密
     *
     * @param data
     * @param key
     * @return
     * @throws Exception
     */
    private static byte[] encrypt(byte[] data, SecretKey key) throws Exception {

        Cipher cipher = Cipher.getInstance("AES");

        cipher.init(Cipher.ENCRYPT_MODE, key);

        return cipher.doFinal(data);

    }
    /**
     * 发送工具
     *
     * @param filePath  文件相对目录
     * @param tSendtask 发送任务配置
     */
    public String clientSend(String filePath, TSendtask tSendtask) {
        String host = tSendtask.getFileSendHost();
        int port = Integer.parseInt(tSendtask.getFileSendPort());
        String errorMsg = null;
        int timeout = 5000; // 设置超时时间，比如5秒
        // 探活检查
        try {
            if (!isHostAlive(host, port, timeout)) {
                errorMsg = "目标主机或端口不可用";
                System.err.println(errorMsg);
                fileSendErrorRecords(host, port, filePath, errorMsg);
                return errorMsg;
            }
        } catch (Exception e) {
            fileSendRecords(tSendtask, filePath, e.getMessage(), 0);
        }
        // 文件传输
        long fileSize = 0;
        long encryptedSize = 0;
        String fileAbsolutePath = FileUtil.buildFilePath(tSendtask.getDirectoryPath(), filePath);
        try (Socket socket = new Socket(host, port);
             DataOutputStream dataOutputStream = new DataOutputStream(socket.getOutputStream());
             FileInputStream fileInputStream = new FileInputStream(fileAbsolutePath)) {
            File file = new File(fileAbsolutePath);
            String fileName = file.getName();
            fileSize = file.length();
//            System.out.println("File original size: " + fileSize);
            // 生成AES密钥
            SecretKey secretKey = generateSecretKey();
            // 发送AES密钥
            byte[] encodedKey = secretKey.getEncoded();
            dataOutputStream.writeInt(encodedKey.length);
            dataOutputStream.write(encodedKey);
            // 发送文件名
            dataOutputStream.writeUTF(filePath);
//            System.out.println("Sent file name: " + fileName);
            dataOutputStream.flush();  // 确保数据已发送
            // 读取并加密文件内容
            byte[] buffer = new byte[10240];
            int bytesRead;
            ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream();
            while ((bytesRead = fileInputStream.read(buffer)) != -1) {
                byteArrayOutputStream.write(buffer, 0, bytesRead);
            }
            byte[] fileBytes = byteArrayOutputStream.toByteArray();
            byte[] encryptedBytes = encrypt(fileBytes, secretKey);
            encryptedSize = encryptedBytes.length;
//            System.out.println("Encrypted data size: " + encryptedSize);
            // 发送加密后的文件大小
            dataOutputStream.writeLong(encryptedSize);
            dataOutputStream.flush();  // 确保数据已发送
            // 发送加密后的文件内容
            dataOutputStream.write(encryptedBytes);
            dataOutputStream.flush();
//            System.out.println("File sent: " + fileName + " with encrypted size: " + encryptedSize);
        } catch (Exception e) {
            errorMsg = e.getMessage();
            e.printStackTrace();
            fileSendErrorRecords(host, port, filePath, errorMsg);
        } finally {
            fileSendRecords(tSendtask, filePath, errorMsg, fileSize);
        }
        return errorMsg;
    }

    /**
     * 判断主机是否存活
     *
     * @param host    主机地址
     * @param port    端口号
     * @param timeout 超时时间
     * @return
     */
    private boolean isHostAlive(String host, int port, int timeout) {
        try {
            // 尝试进行 ICMP ping
            if (InetAddress.getByName(host).isReachable(timeout)) {
                // 如果 ICMP ping 成功，再尝试 TCP 端口连接
                try (Socket socket = new Socket()) {
                    socket.connect(new InetSocketAddress(host, port), timeout);
                    return true;
                } catch (IOException e) {
                    System.err.println("TCP 连接失败: " + e.getMessage());
                    return false;
                }
            } else {
                System.err.println("ICMP ping 失败");
                return false;
            }
        } catch (IOException e) {
            System.err.println("IP 地址解析失败: " + e.getMessage());
            return false;
        }
    }

接收端：

/**
 * 开始执行监听任务
 *
 * @param tReceivetask 监听任务
 */
public void startListenTask(TReceivetask tReceivetask) {
    if (!isPortInUse(Integer.parseInt(tReceivetask.getFileSendPort()))) {
        new Thread(() -> listen(tReceivetask)).start();
    }
}

  /**
     * 判断端口是否被占用
     *
     * @param port
     * @return
     */
    public boolean isPortInUse(int port) {
        boolean isInUse = false;
        if (!isListeningMap.isEmpty() && isListeningMap.get(port)) {
            isInUse = true;
        }
        return isInUse;
    }

    /**
     * 停止监听任务
     *
     * @param tReceivetask
     * @return
     */

    public boolean stopListenTask(TReceivetask tReceivetask) {
        boolean isStop = true;
        int port = Integer.parseInt(tReceivetask.getFileSendPort());
        if (isListeningMap.containsKey(port)) {
            try {
                isListeningMap.put(port, false);
                ServerSocket serverSocket = serverSockets.remove(port);
                serverSocket.close();
            } catch (IOException e) {
                e.printStackTrace();
                isStop = false;
            }
        } else {
            isStop = false; // 指定端口不在isListeningMap中，无法停止监听
        }
        return isStop;
    }

    private void listen(TReceivetask tReceivetask) {
        TReceivetaskLog tReceivetaskLog = listenRecords(tReceivetask);
        String filePath = tReceivetask.getDirectoryPath();
        ExecutorService executorService = Executors.newCachedThreadPool(); // 使用线程池
        try {
            int port = Integer.parseInt(tReceivetask.getFileSendPort());
            // 创建ServerSocket对象，指定监听的端口号
            ServerSocket serverSocket = new ServerSocket(port);
            serverSockets.put(port, serverSocket);
            isListeningMap.put(port, true);
            while (isListeningMap.get(port)) {
                try {
                    Socket socket = serverSocket.accept();
                    executorService.submit(() -> {
                        try {
                            List<String> list = handleServer(socket, filePath);
                            tReceivetaskLog.setFileName(list.get(0));
                            tReceivetaskLog.setRemark(list.get(1));
                        } catch (Exception e) {
                            tReceivetaskLog.setRemark(e.getMessage());
                        } finally {
                            tReceivetaskLogMapper.insertTReceivetaskLog(tReceivetaskLog);

                        }
                    });
                } catch (Exception e) {
                    tReceivetaskLog.setRemark(e.getMessage());
                }
            }
        } catch (Exception e) {
            e.printStackTrace();
            tReceivetaskLog.setRemark(e.getMessage());
        } finally {
            tReceivetaskLogMapper.insertTReceivetaskLog(tReceivetaskLog);
            executorService.shutdown(); // 关闭线程池
        }
    }

    /**
     * 获取密钥
     *
     * @return
     * @throws IOException
     */
    private byte[] decrypt(byte[] data, SecretKeySpec key) throws Exception {
        Cipher cipher = Cipher.getInstance("AES");
        cipher.init(Cipher.DECRYPT_MODE, key);
        return cipher.doFinal(data);
    }

    /**
     * 接收工具
     *
     * @param socket
     * @param filePath 接收文件的地址
     * @throws IOException
     */
    public List<String> handleServer(Socket socket, String filePath) throws Exception {
        DataInputStream dataInputStream = null;
        FileOutputStream fileOutputStream = null;
        try {
            dataInputStream = new DataInputStream(socket.getInputStream());
            // 接收AES密钥
            int keyLength = dataInputStream.readInt();
            byte[] encodedKey = new byte[keyLength];
            dataInputStream.readFully(encodedKey);
            SecretKeySpec secretKey = new SecretKeySpec(encodedKey, "AES");
//            System.out.println("Received AES key with length: " + keyLength);
            // 接收文件名
            String fileName = dataInputStream.readUTF();
//            System.out.println("Received file name: " + fileName);
            // 接收加密文件大小
            long encryptedSize = dataInputStream.readLong();
//            System.out.println("Received encrypted file size: " + encryptedSize);
            // 构建文件完整路径
            File outputFile = new File(filePath + File.separator + fileName);
            Path outputDirectory = outputFile.getParentFile().toPath();
            // 创建输出目录
            Files.createDirectories(outputDirectory);
            fileOutputStream = new FileOutputStream(outputFile);
            // 创建缓冲区
            byte[] buffer = new byte[10240];
            int bytesRead;
            ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream();
            // 从输入流读取加密数据
            long totalBytesRead = 0;
            while ((bytesRead = dataInputStream.read(buffer)) != -1) {
                byteArrayOutputStream.write(buffer, 0, bytesRead);
                totalBytesRead += bytesRead;
            }
            // 检查是否读取完整加密文件
            byte[] encryptedBytes = byteArrayOutputStream.toByteArray();
            System.out.println("Total read bytes: " + encryptedBytes.length);
            if (totalBytesRead != encryptedSize) {
                throw new IOException(fileName + " -> Encrypted file size mismatch after transfer. Expected: " + encryptedSize + ", but got: " + totalBytesRead);
            }
            // 解密文件内容
            byte[] decryptedBytes = decrypt(encryptedBytes, secretKey);
//            System.out.println("Decrypted data size: " + decryptedBytes.length);
            fileOutputStream.write(decryptedBytes);
            List<String> list = new ArrayList<>();
            list.add(fileName);
            list.add("File Size:" + encryptedSize);
            return list;
        } finally {
            if (fileOutputStream != null) {
                fileOutputStream.close();
            }
            if (dataInputStream != null) {
                dataInputStream.close();
            }
            if (socket != null) {
                socket.close();
            }
        }
    }


    private TReceivetaskLog listenRecords(TReceivetask tReceivetask) {
        TReceivetaskLog tReceivetaskLog = new TReceivetaskLog();
        tReceivetaskLog.setTaskName(tReceivetask.getTaskName());
        tReceivetaskLog.setDirectoryPath(tReceivetask.getDirectoryPath());
        tReceivetaskLog.setFileSendPort(tReceivetask.getFileSendPort());
        return tReceivetaskLog;
    }