fix: update search filtering to give more accurate results

Author: adpare

attack-search/__tests__/search-service.test.js

@@ -22,11 +22,17 @@ describe('SearchService', () => {
     });
 
     beforeEach(() => {
+        global.base_url = '/';
         searchService = new SearchService('search-service', null);
+        searchService.render_container = {
+            append: jest.fn(),
+            html: jest.fn(),
+        };
     });
 
     afterEach(async () => {
         searchService = null;
+        delete global.base_url;
     });
 
     it('Access data from mock-index.json', () => {
@@ -111,4 +117,93 @@ describe('SearchService', () => {
         });
 
     });
+
+    test('Keeps only exact ATT&CK ID matches and references, with the object first', async () => {
+        const documents = {
+            1: {
+                id: 1,
+                title: 'TA577, Group G1037',
+                path: '/groups/G1037/index.html',
+                content: 'A group with no reference to the queried technique.',
+                attackId: 'G1037',
+            },
+            2: {
+                id: 2,
+                title: 'Ingress Tool Transfer, Technique T1105 - Enterprise',
+                path: '/techniques/T1105/index.html',
+                content: 'The T1105 technique.',
+                attackId: 'T1105',
+            },
+            3: {
+                id: 3,
+                title: 'A valid reference',
+                path: '/resources/reference/index.html',
+                content: 'This page references T1105.',
+            },
+        };
+        searchService.attackIndex.search = jest.fn().mockResolvedValue([{ field: 'title', result: [1, 2, 3] }]);
+        searchService.resolveSearchResults = jest.fn(async positions => positions.map(position => documents[position]));
+
+        await searchService.query('t1105');
+
+        expect(searchService.allSearchResults.map(result => result.id)).toEqual([2, 3]);
+    });
+
+    test('Treats a four-digit query as an exact ATT&CK ID suffix search', async () => {
+        const documents = {
+            1: {
+                id: 1,
+                title: 'TA577, Group G1037',
+                path: '/groups/G1037/index.html',
+                content: 'A group with no reference to the queried technique.',
+                attackId: 'G1037',
+            },
+            2: {
+                id: 2,
+                title: 'Data from Local System, Technique T1005 - Enterprise',
+                path: '/techniques/T1005/index.html',
+                content: 'The T1005 technique.',
+                attackId: 'T1005',
+            },
+            3: {
+                id: 3,
+                title: 'Matching software, Software S1005',
+                path: '/software/S1005/index.html',
+                content: 'The S1005 software.',
+                attackId: 'S1005',
+            },
+            4: {
+                id: 4,
+                title: 'A valid reference',
+                path: '/resources/reference/index.html',
+                content: 'This page references T1005.',
+            },
+            5: {
+                id: 5,
+                title: 'Data from Local System: Archive Collected Data, Sub-technique T1005.001',
+                path: '/techniques/T1005/001/index.html',
+                content: 'The T1005.001 sub-technique.',
+                attackId: 'T1005.001',
+            },
+        };
+        searchService.attackIndex.search = jest.fn().mockResolvedValue([{ field: 'title', result: [1, 3, 4, 2, 5] }]);
+        searchService.resolveSearchResults = jest.fn(async positions => positions.map(position => documents[position]));
+
+        await searchService.query('1005');
+
+        expect(searchService.allSearchResults.map(result => result.id)).toEqual([2, 5, 3, 4]);
+    });
+
+    test('Preserves result ordering for non-ID queries', async () => {
+        const documents = {
+            1: { id: 1, title: 'First result', path: '/resources/faq/index.html', content: 'Resources' },
+            2: { id: 2, title: 'Second result', path: '/resources/attackcon/index.html', content: 'Resources' },
+        };
+        searchService.attackIndex.search = jest.fn().mockResolvedValue([{ field: 'title', result: [1, 2] }]);
+        searchService.resolveSearchResults = jest.fn(async positions => positions.map(position => documents[position]));
+
+        await searchService.query('Resources');
+
+        expect(searchService.allSearchResults.map(result => result.id)).toEqual([1, 2]);
+    });
 });

attack-search/src/search-service.js

@@ -313,10 +313,84 @@ module.exports = class SearchService {
      *     ]
      */
 
-    this.allSearchResults = await this.#setSearchResults(results);
+    this.allSearchResults = this.#filterAndPromoteExactAttackIdMatches(await this.#setSearchResults(results));
     this.#renderFilteredSearchResults();
   }
 
+  /**
+   * Limits ATT&CK ID searches to matching objects, their sub-techniques, and genuine references.
+   * Non-ID and multi-token queries retain FlexSearch's existing ordering.
+   *
+   * @private
+   * @param {Array<Object>} documents - Search results in their existing relevance order.
+   * @returns {Array<Object>} Exact ATT&CK ID results, with the matching object detail page first when applicable.
+   */
+  #filterAndPromoteExactAttackIdMatches(documents) {
+    const query = this.currentQuery.clean;
+    const isExactAttackId = /^[A-Z]+\d+(?:\.\d+)?$/i.test(query);
+    const isNumericIdSuffix = /^\d{4}$/.test(query);
+    // If user queries for normal text and not attack ids, normal search takes place
+    if (!isExactAttackId && !isNumericIdSuffix) return documents;
+
+    const normalizedQuery = query.toUpperCase();
+
+    // Collect the IDs stored on object-detail search records. Resource and reference pages have no attackId.
+    const candidateAttackIds = documents
+      .map(document => document.attackId?.toUpperCase())
+      .filter(Boolean);
+
+    let directAttackIds;
+    if (isExactAttackId) {
+      // A complete query such as T1005 refers directly to that one ID.
+      directAttackIds = [normalizedQuery];
+    } else {
+      // A numeric query such as 1005 may match T1005, S1005, or another complete ATT&CK ID.
+      const numericSuffixPattern = new RegExp(`^[A-Z]+${normalizedQuery}$`);
+      directAttackIds = [...new Set(candidateAttackIds.filter(attackId => numericSuffixPattern.test(attackId)))];
+    }
+
+    // Include sub-techniques of a matching parent technique, such as T1005.001 for a T1005 query.
+    const subTechniqueIds = candidateAttackIds.filter((attackId) => directAttackIds.some((directAttackId) => (
+      directAttackId.startsWith('T')
+      && !directAttackId.includes('.')
+      && attackId.startsWith(`${directAttackId}.`)
+    )));
+    const matchingAttackIds = [...new Set([...directAttackIds, ...subTechniqueIds])];
+    if (matchingAttackIds.length === 0) return [];
+
+    // Put parent techniques first, then their sub-techniques, followed by other matching ATT&CK object types.
+    const exactMatches = documents.filter(document => matchingAttackIds.includes(document.attackId?.toUpperCase()));
+    exactMatches.sort((first, second) => {
+      const firstIsTechnique = first.attackId.startsWith('T');
+      const secondIsTechnique = second.attackId.startsWith('T');
+      if (firstIsTechnique !== secondIsTechnique) return firstIsTechnique ? -1 : 1;
+
+      const firstIsSubTechnique = first.attackId.includes('.');
+      const secondIsSubTechnique = second.attackId.includes('.');
+      if (firstIsSubTechnique !== secondIsSubTechnique) return firstIsSubTechnique ? 1 : -1;
+
+      return 0;
+    });
+
+    // Escape dots in sub-technique IDs before making one expression that matches only whole IDs.
+    const escapedIds = matchingAttackIds.map(attackId => attackId.replace(/[.*+?^${}()|[\]\\]/g, '\\$&'));
+    // A trailing period is valid sentence punctuation, unless it begins a sub-technique suffix such as .001.
+    const exactIdInText = new RegExp(
+      `(^|[^A-Z0-9.])(?:${escapedIds.join('|')})(?=$|[^A-Z0-9.]|\\.(?!\\d))`,
+      'i',
+    );
+
+    // Add pages that reference a matching ID, but do not add an object-detail page twice.
+    const referencedDocuments = documents.filter((document) => {
+      const title = document.title ?? '';
+      const content = document.content ?? '';
+      const referencesMatchingId = exactIdInText.test(title) || exactIdInText.test(content);
+      return referencesMatchingId && !exactMatches.includes(document);
+    });
+
+    return exactMatches.concat(referencedDocuments);
+  }
+
   /**
    * Renders the search results on the web page based on the given search result page.
    * If the search query is empty, it will show the "Load More Results" button.

modules/search/search.py

@@ -29,6 +29,7 @@
     "x-mitre-data-component": "datacomponents",
     "x-mitre-detection-strategy": "detectionstrategies",
 }
+searchable_object_path_prefixes = set(object_path_prefixes.values()) | {"techniques"}
 
 
 def generate_index():
@@ -64,6 +65,7 @@ def generate_index():
                         "title": title,
                         "path": path,
                         "content": cleancontent,
+                        "attackId": get_search_attack_id(path),
                         "pageType": file_type,
                         "domains": get_domains(title, path, domain_lookup),
                     }
@@ -136,6 +138,19 @@ def should_skip_search_path(path):
     return bool(re.search(r"/sidebar-[^/]+/index\.html$", path))
 
 
+def get_search_attack_id(path):
+    """Return the ATT&CK ID represented by a canonical object detail-page path."""
+    subtechnique_match = re.fullmatch(r"/techniques/(T\d+)/(\d{3})/index\.html", path)
+    if subtechnique_match:
+        return f"{subtechnique_match.group(1)}.{subtechnique_match.group(2)}"
+
+    object_match = re.fullmatch(r"/([^/]+)/([A-Z]+\d+(?:\.\d+)?)/index\.html", path)
+    if object_match and object_match.group(1) in searchable_object_path_prefixes:
+        return object_match.group(2)
+
+    return None
+
+
 def get_domains(title, path=None, domain_lookup=None):
     """Get ATT&CK domains for a search result, preferring source metadata before title inference."""
     if path and domain_lookup and path in domain_lookup: