ReceiveApp/src/main/java/org/example/ReceiveApp.java - master - Dixix404/Using_TLS

ReceiveApp/src/main/java/org/example/ReceiveApp.java
217 строк8 KB
D1x1x
Using TLS
25 дек 2025, 22:33
25 дек 2025, 22:3343fe136
package org.example;

import javax.crypto.Cipher;
import javax.crypto.KeyAgreement;
import javax.crypto.Mac;
import javax.crypto.SecretKey;
import javax.crypto.spec.GCMParameterSpec;
import javax.crypto.spec.SecretKeySpec;
import javax.crypto.spec.DHParameterSpec;

import java.io.*;
import java.math.BigInteger;
import java.net.ServerSocket;
import java.net.Socket;
import java.nio.charset.StandardCharsets;
import java.nio.file.*;
import java.security.*;
import java.security.spec.X509EncodedKeySpec;
import java.util.Arrays;

public class ReceiveApp {

    private static final byte[] MAGIC_DH = new byte[]{'D','H','F','T'};
    private static final byte[] MAGIC_FILE = new byte[]{'E','N','C','F'};
    private static final int VERSION = 1;

    private static final int AES_KEY_BYTES = 32;
    private static final int GCM_TAG_BITS = 128;

    private static final byte REC_UPDATE = 1;
    private static final byte REC_FINAL  = 2;

    public static void main(String[] args) {
        if (args.length != 2) {
            System.err.println("Usage: receive <port> <save-file-path>");
            System.exit(1);
        }

        int port = parsePort(args[0]);
        if (port == -1) {
            System.err.println("Invalid port: " + args[0]);
            System.exit(1);
        }

        Path savePath = Path.of(args[1]);
        if (Files.exists(savePath)) {
            System.err.println("File already exists, will not overwrite: " + savePath);
            System.exit(1);
        }

        Path tmpPath = savePath.resolveSibling(savePath.getFileName() + ".part");

        try (ServerSocket server = new ServerSocket(port);
             Socket socket = server.accept();
             DataInputStream in = new DataInputStream(new BufferedInputStream(socket.getInputStream()));
             DataOutputStream out = new DataOutputStream(new BufferedOutputStream(socket.getOutputStream()))) {

            socket.setTcpNoDelay(true);
            byte[] sharedSecret = dhServerHandshake(in, out);

            byte[] magic = new byte[4];
            in.readFully(magic);
            if (!Arrays.equals(magic, MAGIC_FILE)) throw new IOException("Bad magic (FILE)");

            int ver = in.readInt();
            if (ver != VERSION) throw new IOException("Bad version (FILE): " + ver);

            long expectedSize = in.readLong();
            int ivLen = in.readInt();
            if (ivLen < 12 || ivLen > 32) throw new IOException("Bad IV length: " + ivLen);

            byte[] iv = new byte[ivLen];
            in.readFully(iv);

            SecretKey aesKey = hkdfSha256ToAesKey(sharedSecret, "Using_TLS AES key", AES_KEY_BYTES);

            Cipher dec = Cipher.getInstance("AES/GCM/NoPadding");
            dec.init(Cipher.DECRYPT_MODE, aesKey, new GCMParameterSpec(GCM_TAG_BITS, iv));

            Files.deleteIfExists(tmpPath);

            long written = 0;
            try (OutputStream fileOut = Files.newOutputStream(tmpPath, StandardOpenOption.CREATE_NEW, StandardOpenOption.WRITE)) {
                while (true) {
                    byte recType = in.readByte();
                    int len = in.readInt();
                    if (len < 0 || len > 100_000_000) throw new IOException("Bad record length: " + len);

                    byte[] data = new byte[len];
                    in.readFully(data);

                    if (recType == REC_UPDATE) {
                        byte[] pt = dec.update(data);
                        if (pt != null && pt.length > 0) {
                            fileOut.write(pt);
                            written += pt.length;
                        }
                    } else if (recType == REC_FINAL) {
                        byte[] ptFinal = dec.doFinal(data); // проверка GCM tag
                        if (ptFinal != null && ptFinal.length > 0) {
                            fileOut.write(ptFinal);
                            written += ptFinal.length;
                        }
                        break;
                    } else {
                        throw new IOException("Unknown record type: " + recType);
                    }
                }
            }

            if (written != expectedSize) {
                throw new IOException("Decrypted size mismatch. Expected=" + expectedSize + " got=" + written);
            }

            Files.move(tmpPath, savePath, StandardCopyOption.ATOMIC_MOVE);

            out.writeInt(0); // ACK ok
            out.flush();

        } catch (EOFException e) {
            System.err.println("Connection closed early (EOF). File not fully received.");
            safeDelete(tmpPath);
            System.exit(1);
        } catch (IOException e) {
            System.err.println("I/O error: " + e.getMessage());
            safeDelete(tmpPath);
            System.exit(1);
        } catch (GeneralSecurityException e) {
            System.err.println("Crypto error: " + e.getMessage());
            safeDelete(tmpPath);
            System.exit(1);
        }
    }

    private static byte[] dhServerHandshake(DataInputStream in, DataOutputStream out)
            throws IOException, GeneralSecurityException {

        byte[] magic = new byte[4];
        in.readFully(magic);
        if (!Arrays.equals(magic, MAGIC_DH)) throw new IOException("Bad magic (DH)");

        int ver = in.readInt();
        if (ver != VERSION) throw new IOException("Bad version (DH): " + ver);

        BigInteger p = new BigInteger(readBytes(in));
        BigInteger g = new BigInteger(readBytes(in));
        byte[] clientPubEnc = readBytes(in);

        DHParameterSpec dhSpec = new DHParameterSpec(p, g);

        KeyPairGenerator kpg = KeyPairGenerator.getInstance("DH");
        kpg.initialize(dhSpec);
        KeyPair kp = kpg.generateKeyPair();

        KeyFactory kf = KeyFactory.getInstance("DH");
        PublicKey clientPub = kf.generatePublic(new X509EncodedKeySpec(clientPubEnc));

        KeyAgreement ka = KeyAgreement.getInstance("DH");
        ka.init(kp.getPrivate());
        ka.doPhase(clientPub, true);
        byte[] secret = ka.generateSecret();

        // send server pub
        writeBytes(out, kp.getPublic().getEncoded());
        out.flush();

        return secret;
    }

    private static SecretKey hkdfSha256ToAesKey(byte[] ikm, String info, int keyLenBytes) throws GeneralSecurityException {
        byte[] salt = new byte[32]; // demo salt = 0
        Mac mac = Mac.getInstance("HmacSHA256");

        mac.init(new SecretKeySpec(salt, "HmacSHA256"));
        byte[] prk = mac.doFinal(ikm);

        mac.init(new SecretKeySpec(prk, "HmacSHA256"));
        byte[] t1 = mac.doFinal(concat(new byte[0], info.getBytes(StandardCharsets.UTF_8), new byte[]{1}));

        byte[] okm = Arrays.copyOf(t1, keyLenBytes);
        return new SecretKeySpec(okm, "AES");
    }

    private static byte[] concat(byte[] a, byte[] b, byte[] c) {
        byte[] r = new byte[a.length + b.length + c.length];
        System.arraycopy(a, 0, r, 0, a.length);
        System.arraycopy(b, 0, r, a.length, b.length);
        System.arraycopy(c, 0, r, a.length + b.length, c.length);
        return r;
    }

    private static void writeBytes(DataOutputStream out, byte[] data) throws IOException {
        out.writeInt(data.length);
        out.write(data);
    }

    private static byte[] readBytes(DataInputStream in) throws IOException {
        int len = in.readInt();
        if (len < 0 || len > 50_000_000) throw new IOException("Invalid length: " + len);
        byte[] data = new byte[len];
        in.readFully(data);
        return data;
    }

    private static int parsePort(String s) {
        try {
            int p = Integer.parseInt(s);
            return (p >= 1 && p <= 65535) ? p : -1;
        } catch (NumberFormatException e) {
            return -1;
        }
    }

    private static void safeDelete(Path p) {
        try { Files.deleteIfExists(p); } catch (Exception ignored) {}
    }
}