Merge branch 'unpack-permessage-deflate' into 'master'

Unpack permessage deflate

See merge request packmate/Packmate!3

build.gradle

@@ -19,6 +19,9 @@ configurations {
 
 repositories {
     mavenCentral()
+
+    // удалить после выхода стабильной версии Java-WebSocket
+    maven { url "https://oss.sonatype.org/content/repositories/snapshots" }
 }
 
 dependencies {
@@ -32,6 +35,7 @@ dependencies {
     compile 'org.pcap4j:pcap4j-core:1.8.2'
     compile 'org.pcap4j:pcap4j-packetfactory-static:1.8.2'
     compile group: 'com.google.guava', name: 'guava', version: '28.2-jre'
+    compile group: 'org.java-websocket', name: 'Java-WebSocket', version: '1.5.0-SNAPSHOT'
     compileOnly 'org.projectlombok:lombok'
     runtimeOnly 'org.springframework.boot:spring-boot-devtools'
     runtimeOnly 'org.postgresql:postgresql'

gradle/wrapper/gradle-wrapper.properties

@@ -1,5 +1,5 @@
 distributionBase=GRADLE_USER_HOME
 distributionPath=wrapper/dists
-distributionUrl=https\://services.gradle.org/distributions/gradle-6.2.2-bin.zip
+distributionUrl=https\://services.gradle.org/distributions/gradle-6.3-bin.zip
 zipStoreBase=GRADLE_USER_HOME
 zipStorePath=wrapper/dists

src/main/java/ru/serega6531/packmate/model/CtfService.java

@@ -22,4 +22,6 @@ public class CtfService {
 
     private boolean mergeAdjacentPackets;
 
+    private boolean inflateWebSockets;
+
 }

src/main/java/ru/serega6531/packmate/model/Packet.java

@@ -52,6 +52,18 @@ public class Packet {
 
     private boolean ungzipped;
 
+    private boolean webSocketInflated;
+
     private byte[] content;
 
+    @Transient
+    @JsonIgnore
+    public String getContentString() {
+        return new String(content);
+    }
+
+    public String toString() {
+        return "Packet(id=" + id + ", content=" + getContentString() + ")";
+    }
+
 }

src/main/java/ru/serega6531/packmate/service/StreamOptimizer.java

@@ -24,18 +24,22 @@ import java.util.zip.ZipException;
 public class StreamOptimizer {
 
     private final CtfService service;
-    private final List<Packet> packets;
+    private List<Packet> packets;
 
     private static final byte[] GZIP_HEADER = {0x1f, (byte) 0x8b, 0x08};
 
     /**
      * Вызвать для выполнения оптимизаций на переданном списке пакетов.
      */
-    public void optimizeStream() {
+    public List<Packet> optimizeStream() {
         if (service.isUngzipHttp()) {
             unpackGzip();
         }
 
+        if (service.isInflateWebSockets()) {
+            inflateWebSocket();
+        }
+
         if (service.isUrldecodeHttpRequests()) {
             urldecodeRequests();
         }
@@ -43,6 +47,8 @@ public class StreamOptimizer {
         if (service.isMergeAdjacentPackets()) {
             mergeAdjacentPackets();
         }
+
+        return packets;
     }
 
     /**
@@ -83,6 +89,7 @@ public class StreamOptimizer {
         final List<Packet> cut = packets.subList(start, end);
         final long timestamp = cut.get(0).getTimestamp();
         final boolean ungzipped = cut.stream().anyMatch(Packet::isUngzipped);
+        final boolean webSocketInflated = cut.stream().anyMatch(Packet::isWebSocketInflated);
         boolean incoming = cut.get(0).isIncoming();
         //noinspection OptionalGetWithoutIsPresent
         final byte[] content = cut.stream()
@@ -95,6 +102,7 @@ public class StreamOptimizer {
                 .incoming(incoming)
                 .timestamp(timestamp)
                 .ungzipped(ungzipped)
+                .webSocketInflated(webSocketInflated)
                 .content(content)
                 .build());
     }
@@ -108,7 +116,7 @@ public class StreamOptimizer {
 
         for (Packet packet : packets) {
             if (packet.isIncoming()) {
-                String content = new String(packet.getContent());
+                String content = packet.getContentString();
                 if (content.contains("HTTP/")) {
                     httpStarted = true;
                 }
@@ -145,7 +153,7 @@ public class StreamOptimizer {
                     i = gzipStartPacket + 1;  // продвигаем указатель на следующий после склеенного блок
                 }
             } else if (!packet.isIncoming()) {
-                String content = new String(packet.getContent());
+                String content = packet.getContentString();
 
                 int contentPos = content.indexOf("\r\n\r\n");
                 boolean http = content.startsWith("HTTP/");
@@ -215,6 +223,7 @@ public class StreamOptimizer {
                     .incoming(false)
                     .timestamp(cut.get(0).getTimestamp())
                     .ungzipped(true)
+                    .webSocketInflated(false)
                     .content(newContent)
                     .build();
         } catch (ZipException e) {
@@ -226,4 +235,17 @@ public class StreamOptimizer {
         return null;
     }
 
+    private void inflateWebSocket() {
+        if (!packets.get(0).getContentString().contains("HTTP/")) {
+            return;
+        }
+
+        final WebSocketsParser parser = new WebSocketsParser(packets);
+        if(!parser.isParsed()) {
+            return;
+        }
+
+        packets = parser.getParsedPackets();
+    }
+
 }

src/main/java/ru/serega6531/packmate/service/StreamService.java

@@ -93,7 +93,7 @@ public class StreamService {
 
         countingService.countStream(service.getPort(), packets.size());
 
-        new StreamOptimizer(service, packets).optimizeStream();
+        packets = new StreamOptimizer(service, packets).optimizeStream();
         processUserAgent(packets, stream);
 
         Stream savedStream = save(stream);
@@ -110,7 +110,7 @@ public class StreamService {
     private void processUserAgent(List<Packet> packets, Stream stream) {
         String ua = null;
         for (Packet packet : packets) {
-            String content = new String(packet.getContent());
+            String content = packet.getContentString();
             final Matcher matcher = userAgentPattern.matcher(content);
             if (matcher.find()) {
                 ua = matcher.group(1);

src/main/java/ru/serega6531/packmate/service/WebSocketsParser.java

@@ -0,0 +1,254 @@
+package ru.serega6531.packmate.service;
+
+import lombok.Getter;
+import lombok.extern.slf4j.Slf4j;
+import org.apache.commons.lang3.ArrayUtils;
+import org.java_websocket.drafts.Draft_6455;
+import org.java_websocket.exceptions.InvalidDataException;
+import org.java_websocket.exceptions.InvalidHandshakeException;
+import org.java_websocket.extensions.permessage_deflate.PerMessageDeflateExtension;
+import org.java_websocket.framing.DataFrame;
+import org.java_websocket.framing.Framedata;
+import org.java_websocket.handshake.HandshakeImpl1Client;
+import org.java_websocket.handshake.HandshakeImpl1Server;
+import ru.serega6531.packmate.model.Packet;
+
+import java.nio.ByteBuffer;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.regex.Matcher;
+import java.util.stream.Collectors;
+
+@Slf4j
+public class WebSocketsParser {
+
+    private static final java.util.regex.Pattern WEBSOCKET_KEY_PATTERN =
+            java.util.regex.Pattern.compile("Sec-WebSocket-Key: (.+)\\r\\n");
+    private static final java.util.regex.Pattern WEBSOCKET_EXTENSIONS_PATTERN =
+            java.util.regex.Pattern.compile("Sec-WebSocket-Extensions?: (.+)\\r\\n");
+    private static final java.util.regex.Pattern WEBSOCKET_VERSION_PATTERN =
+            java.util.regex.Pattern.compile("Sec-WebSocket-Version: (\\d+)\\r\\n");
+    private static final java.util.regex.Pattern WEBSOCKET_ACCEPT_PATTERN =
+            java.util.regex.Pattern.compile("Sec-WebSocket-Accept: (.+)\\r\\n");
+
+    private static final String WEBSOCKET_EXTENSION_HEADER = "Sec-WebSocket-Extension: permessage-deflate";
+    private static final String WEBSOCKET_EXTENSIONS_HEADER = "Sec-WebSocket-Extensions: permessage-deflate";
+    private static final String WEBSOCKET_UPGRADE_HEADER = "upgrade: websocket\r\n";
+    private static final String WEBSOCKET_CONNECTION_HEADER = "connection: upgrade\r\n";
+
+    private final List<Packet> packets;
+
+    @Getter
+    private boolean parsed = false;
+    private List<Packet> parsedPackets;
+
+    public WebSocketsParser(List<Packet> packets) {
+        this.packets = packets;
+        detectWebSockets();
+    }
+
+    private void detectWebSockets() {
+        final List<Packet> clientHandshakePackets = packets.stream()
+                .takeWhile(Packet::isIncoming)
+                .collect(Collectors.toList());
+
+        final String clientHandshake = getHandshake(clientHandshakePackets);
+        if (clientHandshake == null) {
+            return;
+        }
+
+        int httpEnd = -1;
+        for (int i = clientHandshakePackets.size(); i < packets.size(); i++) {
+            if (packets.get(i).getContentString().endsWith("\r\n\r\n")) {
+                httpEnd = i + 1;
+                break;
+            }
+        }
+
+        if (httpEnd == -1) {
+            return;
+        }
+
+        final List<Packet> serverHandshakePackets = packets.subList(clientHandshakePackets.size(), httpEnd);
+        final String serverHandshake = getHandshake(serverHandshakePackets);
+        if (serverHandshake == null) {
+            return;
+        }
+
+        HandshakeImpl1Server serverHandshakeImpl = fillServerHandshake(serverHandshake);
+        HandshakeImpl1Client clientHandshakeImpl = fillClientHandshake(clientHandshake);
+
+        if (serverHandshakeImpl == null || clientHandshakeImpl == null) {
+            return;
+        }
+
+        Draft_6455 draft = new Draft_6455(new PerMessageDeflateExtension());
+
+        try {
+            draft.acceptHandshakeAsServer(clientHandshakeImpl);
+            draft.acceptHandshakeAsClient(clientHandshakeImpl, serverHandshakeImpl);
+        } catch (InvalidHandshakeException e) {
+            log.warn("WebSocket handshake", e);
+            return;
+        }
+
+        final List<Packet> wsPackets = packets.subList(
+                httpEnd,
+                packets.size());
+
+        if(wsPackets.isEmpty()) {
+            return;
+        }
+
+        final List<Packet> handshakes = packets.subList(0, httpEnd);
+
+        parse(wsPackets, handshakes, draft);
+        parsed = true;
+    }
+
+    private void parse(final List<Packet> wsPackets, final List<Packet> handshakes, Draft_6455 draft) {
+        List<List<Packet>> sides = sliceToSides(wsPackets);
+        parsedPackets = new ArrayList<>(handshakes);
+
+        for (List<Packet> side : sides) {
+            final Packet lastPacket = side.get(0);
+
+            final byte[] wsContent = side.stream()
+                    .map(Packet::getContent)
+                    .reduce(ArrayUtils::addAll)
+                    .get();
+
+            final ByteBuffer buffer = ByteBuffer.wrap(wsContent);
+            List<Framedata> frames;
+
+            try {
+                frames = draft.translateFrame(buffer);
+            } catch (InvalidDataException e) {
+                log.warn("WebSocket data", e);
+                return;
+            }
+
+            for (Framedata frame : frames) {
+                if(frame instanceof DataFrame) {
+                    parsedPackets.add(Packet.builder()
+                            .content(frame.getPayloadData().array())
+                            .incoming(lastPacket.isIncoming())
+                            .timestamp(lastPacket.getTimestamp())
+                            .ttl(lastPacket.getTtl())
+                            .ungzipped(lastPacket.isUngzipped())
+                            .webSocketInflated(true)
+                            .build()
+                    );
+                }
+            }
+        }
+    }
+
+    public List<Packet> getParsedPackets() {
+        if (!parsed) {
+            throw new IllegalStateException("WS is not parsed");
+        }
+
+        return parsedPackets;
+    }
+
+    private List<List<Packet>> sliceToSides(List<Packet> packets) {
+        List<List<Packet>> result = new ArrayList<>();
+        List<Packet> side = new ArrayList<>();
+        boolean incoming = true;
+
+        for (Packet packet : packets) {
+            if(packet.isIncoming() != incoming) {
+                incoming = packet.isIncoming();
+
+                if(!side.isEmpty()) {
+                    result.add(side);
+                    side = new ArrayList<>();
+                }
+            }
+
+            side.add(packet);
+        }
+
+        if(!side.isEmpty()) {
+            result.add(side);
+        }
+
+        return result;
+    }
+
+    private String getHandshake(final List<Packet> packets) {
+        final String handshake = packets.stream()
+                .map(Packet::getContent)
+                .reduce(ArrayUtils::addAll)
+                .map(String::new)
+                .orElse(null);
+
+        if (handshake == null ||
+                !handshake.toLowerCase().contains(WEBSOCKET_CONNECTION_HEADER) ||
+                !handshake.toLowerCase().contains(WEBSOCKET_UPGRADE_HEADER)) {
+            return null;
+        }
+
+        if (!handshake.contains(WEBSOCKET_EXTENSION_HEADER) &&
+                !handshake.contains(WEBSOCKET_EXTENSIONS_HEADER)) {
+            return null;
+        }
+
+        return handshake;
+    }
+
+    private HandshakeImpl1Client fillClientHandshake(String clientHandshake) {
+        Matcher matcher = WEBSOCKET_VERSION_PATTERN.matcher(clientHandshake);
+        if (!matcher.find()) {
+            return null;
+        }
+        String version = matcher.group(1);
+
+        matcher = WEBSOCKET_KEY_PATTERN.matcher(clientHandshake);
+        if (!matcher.find()) {
+            return null;
+        }
+        String key = matcher.group(1);
+
+        matcher = WEBSOCKET_EXTENSIONS_PATTERN.matcher(clientHandshake);
+        if (!matcher.find()) {
+            return null;
+        }
+        String extensions = matcher.group(1);
+
+        HandshakeImpl1Client clientHandshakeImpl = new HandshakeImpl1Client();
+
+        clientHandshakeImpl.put("Upgrade", "websocket");
+        clientHandshakeImpl.put("Connection", "Upgrade");
+        clientHandshakeImpl.put("Sec-WebSocket-Version", version);
+        clientHandshakeImpl.put("Sec-WebSocket-Key", key);
+        clientHandshakeImpl.put("Sec-WebSocket-Extensions", extensions);
+
+        return clientHandshakeImpl;
+    }
+
+    private HandshakeImpl1Server fillServerHandshake(String serverHandshake) {
+        Matcher matcher = WEBSOCKET_ACCEPT_PATTERN.matcher(serverHandshake);
+        if (!matcher.find()) {
+            return null;
+        }
+        String accept = matcher.group(1);
+
+        matcher = WEBSOCKET_EXTENSIONS_PATTERN.matcher(serverHandshake);
+        if (!matcher.find()) {
+            return null;
+        }
+        String extensions = matcher.group(1);
+
+        HandshakeImpl1Server serverHandshakeImpl = new HandshakeImpl1Server();
+
+        serverHandshakeImpl.put("Upgrade", "websocket");
+        serverHandshakeImpl.put("Connection", "Upgrade");
+        serverHandshakeImpl.put("Sec-WebSocket-Accept", accept);
+        serverHandshakeImpl.put("Sec-WebSocket-Extensions", extensions);
+
+        return serverHandshakeImpl;
+    }
+
+}

src/test/java/ru/serega6531/packmate/StreamOptimizerTest.java

@@ -28,8 +28,8 @@ class StreamOptimizerTest {
         List<Packet> list = new ArrayList<>();
         list.add(p);
 
-        new StreamOptimizer(service, list).optimizeStream();
+        list = new StreamOptimizer(service, list).optimizeStream();
-        final String processed = new String(list.get(0).getContent());
+        final String processed = list.get(0).getContentString();
         assertTrue(processed.contains("aaabbb"));
     }
 
@@ -42,8 +42,8 @@ class StreamOptimizerTest {
         List<Packet> list = new ArrayList<>();
         list.add(p);
 
-        new StreamOptimizer(service, list).optimizeStream();
+        list = new StreamOptimizer(service, list).optimizeStream();
-        final String processed = new String(list.get(0).getContent());
+        final String processed = list.get(0).getContentString();
         assertTrue(processed.contains("а б"));
     }
 
@@ -67,7 +67,7 @@ class StreamOptimizerTest {
         list.add(p5);
         list.add(p6);
 
-        new StreamOptimizer(service, list).optimizeStream();
+        list = new StreamOptimizer(service, list).optimizeStream();
 
         assertEquals(4, list.size());
         assertEquals(2, list.get(1).getContent().length);