CustomSignatureService.java

package com.tracepcap.analysis.service;

import java.io.File;
import java.io.FileInputStream;
import java.io.IOException;
import java.net.InetAddress;
import java.net.UnknownHostException;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.regex.Pattern;
import java.util.regex.PatternSyntaxException;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Service;
import org.yaml.snakeyaml.LoaderOptions;
import org.yaml.snakeyaml.Yaml;
import org.yaml.snakeyaml.constructor.SafeConstructor;

/**
 * Loads custom detection rules from {@code signatures.yml} and applies them to each conversation
 * after nDPI enrichment. Matched rule names are appended to the conversation's {@code
 * customSignatures} list, making them visible in the UI alongside nDPI's built-in risk flags.
 *
 * <p>The file is reloaded on every analysis run so admins can update rules without restarting the
 * application.
 *
 * <p>Supported match fields (all optional; a rule fires when ALL specified fields match):
 *
 * <ul>
 *   <li>{@code ip} — exact match against srcIp OR dstIp
 *   <li>{@code cidr} — CIDR range match against srcIp OR dstIp (e.g. {@code 10.0.0.0/8})
 *   <li>{@code srcPort} — exact match against source port
 *   <li>{@code dstPort} — exact match against destination port
 *   <li>{@code ja3} — exact match against ja3Client OR ja3Server
 *   <li>{@code hostname}— exact or wildcard match against SNI hostname (e.g. {@code *.evil.com})
 *   <li>{@code app} — case-insensitive match against nDPI appName
 *   <li>{@code protocol} — case-insensitive match against transport/network protocol (e.g. TCP,
 *       UDP, ICMP)
 * </ul>
 *
 * <p>Top-level rule keys (alongside {@code match}):
 *
 * <ul>
 *   <li>{@code payload_contains} — list of {@code {ascii: "..."}} or {@code {hex: "..."}} entries;
 *       OR-matched by default, AND-matched when {@code match_all: true} is set on the rule.
 *   <li>{@code payload_regex} — list of {@code {pattern: "..."}} entries (standard Java regex
 *       applied against the ASCII/UTF-8 decoded payload); supports optional {@code
 *       case_insensitive: true} per entry. Uses the same {@code match_all} flag as {@code
 *       payload_contains}. Payloads are capped at {@value #MAX_REGEX_PAYLOAD_BYTES} bytes per
 *       packet to prevent catastrophic backtracking.
 * </ul>
 */
@Slf4j
@Service
public class CustomSignatureService {

  static final int MAX_REGEX_PAYLOAD_BYTES = 65536;

  private final Map<String, Pattern> patternCache = new ConcurrentHashMap<>();

  @Value("${tracepcap.signatures.path:/app/config/signatures.yml}")
  private String signaturesPath;

  /**
   * Evaluates all loaded rules against each conversation and appends matched rule names to the
   * conversation's {@code flowRisks} list in-place.
   */
  public void applySignatures(List<PcapParserService.ConversationInfo> conversations) {
    List<Map<String, Object>> rules = loadRules();
    if (rules.isEmpty() || conversations.isEmpty()) return;

    int matchCount = 0;
    for (PcapParserService.ConversationInfo conv : conversations) {
      for (Map<String, Object> rule : rules) {
        String name = (String) rule.get("name");
        if (name == null || name.isBlank()) continue;

        @SuppressWarnings("unchecked")
        Map<String, Object> match = (Map<String, Object>) rule.get("match");

        @SuppressWarnings("unchecked")
        List<Map<String, Object>> payloadContains =
            (List<Map<String, Object>>) rule.get("payload_contains");

        Object payloadRegexObj = rule.get("payload_regex");
        @SuppressWarnings("unchecked")
        List<Map<String, Object>> payloadRegex =
            payloadRegexObj instanceof List ? (List<Map<String, Object>>) payloadRegexObj : null;

        // A rule must specify at least one criterion
        boolean hasMatch = match != null && !match.isEmpty();
        boolean hasPayload = payloadContains != null && !payloadContains.isEmpty();
        boolean hasRegex = payloadRegex != null && !payloadRegex.isEmpty();
        if (!hasMatch && !hasPayload && !hasRegex) continue;

        // All specified criteria must pass (AND between match block, payload_contains, payload_regex)
        if (hasMatch && !matches(conv, match)) continue;

        boolean matchAll = Boolean.TRUE.equals(rule.get("match_all"));

        if (hasPayload) {
          if (!payloadContainsMatch(conv.getPackets(), payloadContains, matchAll)) continue;
        }

        if (hasRegex) {
          if (!payloadRegexMatch(conv.getPackets(), payloadRegex, matchAll)) continue;
        }

        if (!conv.getCustomSignatures().contains(name)) {
          conv.getCustomSignatures().add(name);
          matchCount++;
        }
      }
    }
    if (matchCount > 0) {
      log.info(
          "Custom signatures: {} match(es) across {} conversations",
          matchCount,
          conversations.size());
    }
  }

  /**
   * Returns a map of IP address → custom device type for any IPs that were involved in
   * conversations matched by a rule containing a {@code device_type} field. Call this after {@link
   * #applySignatures} so the conversations already carry their matched rule names.
   *
   * <p>Example YAML:
   *
   * <pre>
   * - name: my_cctv
   *   device_type: "CCTV Camera"
   *   match:
   *     ip: "192.168.1.50"
   * </pre>
   */
  public Map<String, String> getDeviceTypeOverrides(
      List<PcapParserService.ConversationInfo> conversations) {
    List<Map<String, Object>> rules = loadRules();
    if (rules.isEmpty()) return Map.of();

    // Build rule-name → device_type map in one pass over rules
    Map<String, String> ruleDeviceType = new HashMap<>();
    for (Map<String, Object> rule : rules) {
      String deviceType = (String) rule.get("device_type");
      String name = (String) rule.get("name");
      if (deviceType != null && !deviceType.isBlank() && name != null && !name.isBlank()) {
        ruleDeviceType.put(name, deviceType);
      }
    }
    if (ruleDeviceType.isEmpty()) return Map.of();

    // One pass over conversations — look up each matched signature in the map
    Map<String, String> overrides = new HashMap<>();
    for (PcapParserService.ConversationInfo conv : conversations) {
      for (String sig : conv.getCustomSignatures()) {
        String deviceType = ruleDeviceType.get(sig);
        if (deviceType != null) {
          if (conv.getSrcIp() != null) overrides.putIfAbsent(conv.getSrcIp(), deviceType);
          if (conv.getDstIp() != null) overrides.putIfAbsent(conv.getDstIp(), deviceType);
        }
      }
    }
    return overrides;
  }

  // -------------------------------------------------------------------------
  // Internal helpers
  // -------------------------------------------------------------------------

  /**
   * Load and parse the signatures file. Returns an empty list if the file is missing or invalid.
   */
  @SuppressWarnings("unchecked")
  private List<Map<String, Object>> loadRules() {
    File file = new File(signaturesPath);
    if (!file.exists()) {
      return List.of();
    }

    try (FileInputStream fis = new FileInputStream(file)) {
      Yaml yaml = new Yaml(new SafeConstructor(new LoaderOptions()));
      Map<String, Object> root = yaml.load(fis);
      if (root == null || !root.containsKey("signatures")) return List.of();

      Object signaturesObj = root.get("signatures");
      if (!(signaturesObj instanceof List)) return List.of();

      List<?> raw = (List<?>) signaturesObj;
      List<Map<String, Object>> rules = new ArrayList<>();
      for (Object item : raw) {
        if (item instanceof Map) {
          rules.add((Map<String, Object>) item);
        }
      }
      return rules;
    } catch (IOException e) {
      log.warn("Could not read signatures file at {}: {}", signaturesPath, e.getMessage());
      return List.of();
    } catch (Exception e) {
      log.warn("Failed to parse signatures file at {}: {}", signaturesPath, e.getMessage());
      return List.of();
    }
  }

  /** Returns true if the conversation satisfies ALL non-null criteria in the match block. */
  private boolean matches(PcapParserService.ConversationInfo conv, Map<String, Object> match) {
    // ip — exact match against srcIp OR dstIp
    if (match.containsKey("ip")) {
      String ip = String.valueOf(match.get("ip"));
      if (!ip.equals(conv.getSrcIp()) && !ip.equals(conv.getDstIp())) return false;
    }

    // cidr — CIDR range match against srcIp OR dstIp
    if (match.containsKey("cidr")) {
      String cidr = String.valueOf(match.get("cidr"));
      if (!inCidr(conv.getSrcIp(), cidr) && !inCidr(conv.getDstIp(), cidr)) return false;
    }

    // srcPort — exact match
    if (match.containsKey("srcPort")) {
      int port = toInt(match.get("srcPort"));
      if (!Integer.valueOf(port).equals(conv.getSrcPort())) return false;
    }

    // dstPort — exact match
    if (match.containsKey("dstPort")) {
      int port = toInt(match.get("dstPort"));
      if (!Integer.valueOf(port).equals(conv.getDstPort())) return false;
    }

    // ja3 — exact match against ja3Client OR ja3Server
    if (match.containsKey("ja3")) {
      String ja3 = String.valueOf(match.get("ja3"));
      boolean clientMatch = ja3.equals(conv.getJa3Client());
      boolean serverMatch = ja3.equals(conv.getJa3Server());
      if (!clientMatch && !serverMatch) return false;
    }

    // hostname — exact or wildcard match against SNI hostname
    if (match.containsKey("hostname")) {
      String pattern = String.valueOf(match.get("hostname"));
      if (!hostnameMatches(conv.getHostname(), pattern)) return false;
    }

    // app — case-insensitive match against nDPI appName
    if (match.containsKey("app")) {
      String app = String.valueOf(match.get("app"));
      if (conv.getAppName() == null || !app.equalsIgnoreCase(conv.getAppName())) return false;
    }

    // protocol — case-insensitive match against transport/network protocol (e.g. TCP, UDP, ICMP)
    if (match.containsKey("protocol")) {
      String protocol = String.valueOf(match.get("protocol"));
      if (conv.getProtocol() == null || !protocol.equalsIgnoreCase(conv.getProtocol()))
        return false;
    }

    return true;
  }

  /** Returns true if {@code ip} falls within the given CIDR range. */
  private boolean inCidr(String ip, String cidr) {
    if (ip == null || cidr == null) return false;
    try {
      String[] parts = cidr.split("/");
      if (parts.length != 2) return false;
      int prefixLen = Integer.parseInt(parts[1]);

      InetAddress network = InetAddress.getByName(parts[0]);
      InetAddress address = InetAddress.getByName(ip);

      // Must be the same address family
      byte[] netBytes = network.getAddress();
      byte[] addrBytes = address.getAddress();
      if (netBytes.length != addrBytes.length) return false;

      // Compare the first prefixLen bits
      int fullBytes = prefixLen / 8;
      int remainingBits = prefixLen % 8;

      for (int i = 0; i < fullBytes; i++) {
        if (netBytes[i] != addrBytes[i]) return false;
      }
      if (remainingBits > 0 && fullBytes < netBytes.length) {
        int mask = 0xFF & (0xFF << (8 - remainingBits));
        if ((netBytes[fullBytes] & mask) != (addrBytes[fullBytes] & mask)) return false;
      }
      return true;
    } catch (UnknownHostException | NumberFormatException e) {
      log.warn("Invalid CIDR '{}': {}", cidr, e.getMessage());
      return false;
    }
  }

  /**
   * Matches a hostname against a pattern. Supports a leading wildcard: {@code *.example.com}
   * matches any subdomain at any depth (e.g. {@code foo.example.com}, {@code foo.bar.example.com})
   * as well as the apex domain itself ({@code example.com}).
   */
  private boolean hostnameMatches(String hostname, String pattern) {
    if (hostname == null || pattern == null) return false;
    if (!pattern.startsWith("*.")) {
      return pattern.equalsIgnoreCase(hostname);
    }
    // Wildcard: strip the "*." and check that hostname ends with the remainder,
    // preceded by a "." or equal to it (covers the apex domain itself).
    // Matches any depth: "*.example.com" matches "a.example.com", "a.b.example.com", etc.
    String suffix = pattern.substring(2).toLowerCase();
    String h = hostname.toLowerCase();
    return h.equals(suffix) || h.endsWith("." + suffix);
  }

  private int toInt(Object value) {
    if (value instanceof Number) return ((Number) value).intValue();
    return Integer.parseInt(value.toString());
  }

  /**
   * Returns true if at least one (OR) or all (AND when {@code matchAll=true}) of the given payload
   * patterns are found in the raw payload bytes of any packet in the conversation.
   *
   * <p>Each pattern entry is a map with one key: {@code ascii} (plain text) or {@code hex} (hex
   * string, with optional {@code 0x} prefix and space/colon separators). Patterns are matched
   * against the packet's lowercase hex payload string.
   */
  private boolean payloadContainsMatch(
      List<PcapParserService.PacketInfo> packets,
      List<Map<String, Object>> patterns,
      boolean matchAll) {
    for (Map<String, Object> pattern : patterns) {
      String hexNeedle = null;
      if (pattern.containsKey("ascii")) {
        hexNeedle = asciiToHex(String.valueOf(pattern.get("ascii")));
      } else if (pattern.containsKey("hex")) {
        hexNeedle = normalizeHex(String.valueOf(pattern.get("hex")));
      }
      if (hexNeedle == null || hexNeedle.isEmpty()) {
        if (matchAll) return false; // empty/invalid pattern counts as unmatched in AND mode
        continue;
      }

      final String needle = hexNeedle;
      boolean found =
          packets.stream().anyMatch(p -> p.getPayload() != null && p.getPayload().contains(needle));

      if (matchAll && !found) return false;
      if (!matchAll && found) return true;
    }
    // matchAll: all patterns passed. OR: none matched.
    return matchAll;
  }

  /** Converts an ASCII string to its lowercase hex representation (e.g. "GET" → "474554"). */
  private String asciiToHex(String ascii) {
    if (ascii == null) return "";
    StringBuilder sb = new StringBuilder(ascii.length() * 2);
    for (char c : ascii.toCharArray()) {
      sb.append(String.format("%02x", (int) c));
    }
    return sb.toString();
  }

  /**
   * Strips hex formatting ({@code 0x} prefix, spaces, colons, hyphens) and lowercases the result.
   */
  private String normalizeHex(String hex) {
    if (hex == null) return "";
    String s = hex.toLowerCase();
    if (s.startsWith("0x")) s = s.substring(2);
    return s.replaceAll("[\\s:\\-]", "");
  }

  /**
   * Returns true if at least one (OR) or all (AND when {@code matchAll=true}) of the given regex
   * patterns match any packet payload in the conversation.
   *
   * <p>Each entry is a map with a required {@code pattern} key (standard Java regex) and an
   * optional {@code case_insensitive: true} flag. Payloads are decoded from hex to ASCII/UTF-8 and
   * capped at {@value #MAX_REGEX_PAYLOAD_BYTES} bytes per packet to prevent backtracking issues.
   */
  private boolean payloadRegexMatch(
      List<PcapParserService.PacketInfo> packets,
      List<Map<String, Object>> patterns,
      boolean matchAll) {
    // Lazily decoded payloads — populated on first access for each packet index
    String[] decoded = new String[packets.size()];
    for (Map<String, Object> entry : patterns) {
      Object patternObj = entry.get("pattern");
      if (patternObj == null) {
        if (matchAll) return false;
        continue;
      }
      String patternStr = patternObj.toString();
      if (patternStr.isBlank()) {
        if (matchAll) return false;
        continue;
      }

      final int flags = Boolean.TRUE.equals(entry.get("case_insensitive"))
          ? Pattern.CASE_INSENSITIVE | Pattern.UNICODE_CASE
          : 0;

      final String cacheKey = patternStr + "::" + flags;
      Pattern compiled = patternCache.computeIfAbsent(cacheKey, k -> {
        try {
          return Pattern.compile(patternStr, flags);
        } catch (PatternSyntaxException e) {
          log.warn("Invalid regex pattern '{}': {}", patternStr, e.getMessage());
          return null;
        }
      });

      if (compiled == null) {
        if (matchAll) return false;
        continue;
      }

      boolean found = false;
      for (int i = 0; i < packets.size(); i++) {
        String payloadHex = packets.get(i).getPayload();
        if (payloadHex == null || payloadHex.isEmpty()) continue;
        if (decoded[i] == null) decoded[i] = hexToAscii(payloadHex);
        if (compiled.matcher(decoded[i]).find()) {
          found = true;
          break;
        }
      }

      if (matchAll && !found) return false;
      if (!matchAll && found) return true;
    }
    return matchAll;
  }

  /** Decodes a hex-encoded payload to a UTF-8 string, capped at {@value #MAX_REGEX_PAYLOAD_BYTES} bytes. */
  private String hexToAscii(String hex) {
    if (hex == null) return "";
    int hexLen = Math.min(hex.length() & ~1, MAX_REGEX_PAYLOAD_BYTES * 2);
    byte[] bytes = new byte[hexLen / 2];
    int out = 0;
    for (int i = 0; i < hexLen; i += 2) {
      int h1 = Character.digit(hex.charAt(i), 16);
      int h2 = Character.digit(hex.charAt(i + 1), 16);
      bytes[out++] = (h1 == -1 || h2 == -1) ? (byte) '?' : (byte) ((h1 << 4) | h2);
    }
    return new String(bytes, 0, out, StandardCharsets.UTF_8);
  }
}