diff --git a/src/main/java/ru/serega6531/packmate/PcapWorker.java b/src/main/java/ru/serega6531/packmate/PcapWorker.java index 6cc9b1f..a394c28 100644 --- a/src/main/java/ru/serega6531/packmate/PcapWorker.java +++ b/src/main/java/ru/serega6531/packmate/PcapWorker.java @@ -14,9 +14,9 @@ import org.springframework.beans.factory.annotation.Autowired; import org.springframework.beans.factory.annotation.Value; import org.springframework.stereotype.Component; import ru.serega6531.packmate.model.CtfService; -import ru.serega6531.packmate.model.Stream; import ru.serega6531.packmate.model.UnfinishedStream; -import ru.serega6531.packmate.service.*; +import ru.serega6531.packmate.service.ServicesService; +import ru.serega6531.packmate.service.StreamService; import javax.annotation.PreDestroy; import java.net.Inet4Address; @@ -31,8 +31,6 @@ public class PcapWorker { private final ServicesService servicesService; private final StreamService streamService; - private final PacketService packetService; - private final StreamSubscriptionService subscriptionService; private final PcapNetworkInterface device; private PcapHandle pcap = null; @@ -51,14 +49,10 @@ public class PcapWorker { @Autowired public PcapWorker(ServicesService servicesService, StreamService streamService, - PacketService packetService, - StreamSubscriptionService subscriptionService, @Value("${interface-name}") String interfaceName, @Value("${local-ip}") String localIp) throws PcapNativeException { this.servicesService = servicesService; this.streamService = streamService; - this.packetService = packetService; - this.subscriptionService = subscriptionService; this.localIp = localIp; @@ -115,7 +109,7 @@ public class PcapWorker { boolean fin = false; boolean rst = false; - if(rawPacket.contains(IpV4Packet.class)){ + if (rawPacket.contains(IpV4Packet.class)) { final IpV4Packet.IpV4Header header = rawPacket.get(IpV4Packet.class).getHeader(); sourceIp = header.getSrcAddr(); destIp = header.getDstAddr(); @@ -123,7 +117,7 @@ public class PcapWorker { destIpString = header.getDstAddr().getHostAddress(); } - if(rawPacket.contains(TcpPacket.class)) { + if (rawPacket.contains(TcpPacket.class)) { final TcpPacket packet = rawPacket.get(TcpPacket.class); final TcpPacket.TcpHeader header = packet.getHeader(); sourcePort = header.getSrcPort().valueAsInt(); @@ -133,7 +127,7 @@ public class PcapWorker { rst = header.getRst(); content = packet.getRawData(); protocol = Protocol.TCP; - } else if(rawPacket.contains(UdpPacket.class)) { + } else if (rawPacket.contains(UdpPacket.class)) { final UdpPacket packet = rawPacket.get(UdpPacket.class); final UdpPacket.UdpHeader header = packet.getHeader(); sourcePort = header.getSrcPort().valueAsInt(); @@ -142,10 +136,11 @@ public class PcapWorker { protocol = Protocol.UDP; } - if(sourceIpString != null && sourcePort != -1) { - final Optional serviceOptional = findService(sourceIpString, sourcePort, destIpString, destPort); + if (sourceIpString != null && sourcePort != -1) { + final Optional serviceOptional = + servicesService.findService(localIp, sourceIpString, sourcePort, destIpString, destPort); - if(serviceOptional.isPresent()) { + if (serviceOptional.isPresent()) { String sourceIpAndPort = sourceIpString + ":" + sourcePort; String destIpAndPort = destIpString + ":" + destPort; @@ -157,7 +152,7 @@ public class PcapWorker { .content(content) .build(); - if(unfinishedStreams.containsKey(stream)) { + if (unfinishedStreams.containsKey(stream)) { unfinishedStreams.get(stream).add(packet); } else { List packets = new ArrayList<>(); @@ -169,69 +164,35 @@ public class PcapWorker { protocol.name().toLowerCase(), serviceOptional.get(), sourceIpString, sourcePort, destIpString, destPort, unfinishedStreams.get(stream).size()); - if(protocol == Protocol.TCP) { - if(!fins.containsKey(stream)) { + if (protocol == Protocol.TCP) { + if (!fins.containsKey(stream)) { fins.put(stream, new HashSet<>()); } - if(!acks.containsKey(stream)) { + if (!acks.containsKey(stream)) { acks.put(stream, new HashSet<>()); } final Set finsForStream = fins.get(stream); final Set acksForStream = acks.get(stream); - if(fin) { + if (fin) { finsForStream.add(sourceIpAndPort); } - if(ack && finsForStream.contains(destIpAndPort)) { // проверяем destIp, потому что ищем ответ на его fin + if (ack && finsForStream.contains(destIpAndPort)) { // проверяем destIp, потому что ищем ответ на его fin acksForStream.add(sourceIpAndPort); } - if(rst || (acksForStream.contains(sourceIpAndPort) && acksForStream.contains(destIpAndPort))) { - log.info("Конец стрима"); - saveStream(stream); + if (rst || (acksForStream.contains(sourceIpAndPort) && acksForStream.contains(destIpAndPort))) { + streamService.saveNewStream(stream, unfinishedStreams.get(stream)); + + unfinishedStreams.remove(stream); + fins.remove(stream); + acks.remove(stream); } } } } } - - private void saveStream(UnfinishedStream unfinishedStream) { - final List packets = unfinishedStreams.get(unfinishedStream); - - Stream stream = new Stream(); - stream.setProtocol(unfinishedStream.getProtocol()); - stream.setStartTimestamp(packets.get(0).getTimestamp()); - stream.setEndTimestamp(packets.get(packets.size() - 1).getTimestamp()); - stream.setService(findService( - unfinishedStream.getFirstIp().getHostAddress(), - unfinishedStream.getFirstPort(), - unfinishedStream.getSecondIp().getHostAddress(), - unfinishedStream.getSecondPort() - ).get()); - - Stream savedStream = streamService.save(stream); - - List savedPackets = new ArrayList<>(); - for (ru.serega6531.packmate.model.Packet packet : packets) { - packet.setStream(savedStream); - savedPackets.add(packetService.save(packet)); - } - - savedStream.setPackets(savedPackets); - savedStream = streamService.save(savedStream); - subscriptionService.broadcastNewStream(savedStream); - } - - private Optional findService(String firstIp, int firstPort, String secondIp, int secondPort) { - if(firstIp.equals(localIp)) { - return servicesService.findByPort(firstPort); - } else if(secondIp.equals(localIp)) { - return servicesService.findByPort(secondPort); - } - - return Optional.empty(); - } } diff --git a/src/main/java/ru/serega6531/packmate/model/CtfService.java b/src/main/java/ru/serega6531/packmate/model/CtfService.java index c34ec8b..574b896 100644 --- a/src/main/java/ru/serega6531/packmate/model/CtfService.java +++ b/src/main/java/ru/serega6531/packmate/model/CtfService.java @@ -10,7 +10,7 @@ import java.util.List; @Data @Entity @Table(name = "service") -@ToString(of = {"port", "name"}) +@ToString(exclude = "streams") public class CtfService { @Id diff --git a/src/main/java/ru/serega6531/packmate/model/Pattern.java b/src/main/java/ru/serega6531/packmate/model/Pattern.java index 2692f6b..11d01b7 100644 --- a/src/main/java/ru/serega6531/packmate/model/Pattern.java +++ b/src/main/java/ru/serega6531/packmate/model/Pattern.java @@ -3,9 +3,8 @@ package ru.serega6531.packmate.model; import lombok.Data; import org.hibernate.annotations.GenericGenerator; -import javax.persistence.Entity; -import javax.persistence.GeneratedValue; -import javax.persistence.Id; +import javax.persistence.*; +import java.util.List; @Data @Entity @@ -30,4 +29,7 @@ public class Pattern { private String color; // для вставки в css + @ManyToMany(mappedBy = "foundPatterns", cascade = CascadeType.ALL) + private List matchedStreams; + } diff --git a/src/main/java/ru/serega6531/packmate/model/Stream.java b/src/main/java/ru/serega6531/packmate/model/Stream.java index 98bfd33..b4efb63 100644 --- a/src/main/java/ru/serega6531/packmate/model/Stream.java +++ b/src/main/java/ru/serega6531/packmate/model/Stream.java @@ -2,14 +2,15 @@ package ru.serega6531.packmate.model; import com.fasterxml.jackson.annotation.JsonIgnore; import lombok.Data; +import lombok.ToString; import org.hibernate.annotations.GenericGenerator; import ru.serega6531.packmate.Protocol; import javax.persistence.*; import java.util.List; -import java.util.Set; @Data +@ToString(exclude = "packets") @Entity @GenericGenerator( name = "stream_generator", @@ -40,6 +41,7 @@ public class Stream { private long endTimestamp; - private Set foundPatterns; + @ManyToMany(cascade = CascadeType.ALL) + private List foundPatterns; } diff --git a/src/main/java/ru/serega6531/packmate/service/ServicesService.java b/src/main/java/ru/serega6531/packmate/service/ServicesService.java index 3da4430..7cf578f 100644 --- a/src/main/java/ru/serega6531/packmate/service/ServicesService.java +++ b/src/main/java/ru/serega6531/packmate/service/ServicesService.java @@ -20,6 +20,16 @@ public class ServicesService { this.repository = repository; } + public Optional findService(String localIp, String firstIp, int firstPort, String secondIp, int secondPort) { + if(firstIp.equals(localIp)) { + return findByPort(firstPort); + } else if(secondIp.equals(localIp)) { + return findByPort(secondPort); + } + + return Optional.empty(); + } + public Optional findByPort(int port) { return repository.findById(port); } diff --git a/src/main/java/ru/serega6531/packmate/service/StreamService.java b/src/main/java/ru/serega6531/packmate/service/StreamService.java index b289f95..0fa6c39 100644 --- a/src/main/java/ru/serega6531/packmate/service/StreamService.java +++ b/src/main/java/ru/serega6531/packmate/service/StreamService.java @@ -2,15 +2,16 @@ package ru.serega6531.packmate.service; import lombok.extern.slf4j.Slf4j; import org.springframework.beans.factory.annotation.Autowired; +import org.springframework.beans.factory.annotation.Value; import org.springframework.stereotype.Service; +import org.springframework.transaction.annotation.Transactional; +import ru.serega6531.packmate.model.Packet; import ru.serega6531.packmate.model.Pattern; import ru.serega6531.packmate.model.Stream; +import ru.serega6531.packmate.model.UnfinishedStream; import ru.serega6531.packmate.repository.StreamRepository; -import java.util.HashSet; -import java.util.List; -import java.util.Optional; -import java.util.Set; +import java.util.*; @Service @Slf4j @@ -18,26 +19,68 @@ public class StreamService { private final StreamRepository repository; private final PatternService patternService; + private final ServicesService servicesService; + private final PacketService packetService; + private final StreamSubscriptionService subscriptionService; + + private final String localIp; @Autowired - public StreamService(StreamRepository repository, PatternService patternService) { + public StreamService(StreamRepository repository, + PatternService patternService, + ServicesService servicesService, + PacketService packetService, + StreamSubscriptionService subscriptionService, + @Value("${local-ip}") String localIp) { this.repository = repository; this.patternService = patternService; + this.servicesService = servicesService; + this.packetService = packetService; + this.subscriptionService = subscriptionService; + this.localIp = localIp; + } + + @Transactional + public void saveNewStream(UnfinishedStream unfinishedStream, List packets) { + Stream stream = new Stream(); + stream.setProtocol(unfinishedStream.getProtocol()); + stream.setStartTimestamp(packets.get(0).getTimestamp()); + stream.setEndTimestamp(packets.get(packets.size() - 1).getTimestamp()); + stream.setService(servicesService.findService( + localIp, + unfinishedStream.getFirstIp().getHostAddress(), + unfinishedStream.getFirstPort(), + unfinishedStream.getSecondIp().getHostAddress(), + unfinishedStream.getSecondPort() + ).get()); + + Stream savedStream = save(stream); + + List savedPackets = new ArrayList<>(); + Set matches = new HashSet<>(); + + for (ru.serega6531.packmate.model.Packet packet : packets) { + packet.setStream(savedStream); + savedPackets.add(packetService.save(packet)); + matches.addAll(patternService.findMatching(packet.getContent())); + } + + savedStream.setFoundPatterns(new ArrayList<>(matches)); + savedStream.setPackets(savedPackets); + savedStream = save(savedStream); + + subscriptionService.broadcastNewStream(savedStream); } public Stream save(Stream stream) { - if(!stream.getPackets().isEmpty()) { - Set matches = new HashSet<>(); - - stream.getPackets().forEach(packet -> { - matches.addAll(patternService.findMatching(packet.getContent())); - }); - - stream.setFoundPatterns(matches); + Stream saved; + if(stream.getId() == null) { + saved = repository.save(stream); + log.info("Создан стрим с id {}", saved.getId()); + } else { + saved = repository.save(stream); } - final Stream saved = repository.save(stream); - log.info("Создан стрим с id {}", saved.getId()); return saved; }