[8.x] [EDR Workflows] Add support to multiple agent id's in a po…

…licy for Osquery (#193987) (#194138)

# Backport

This will backport the following commits from `main` to `8.x`:
- [[EDR Workflows] Add support to multiple agent id's in a policy
for Osquery (#193987)](#193987)

<!--- Backport version: 9.4.3 -->

### Questions ?
Please refer to the [Backport tool
documentation](https://github.com/sqren/backport)

<!--BACKPORT [{"author":{"name":"Tomasz
Ciecierski","email":"[email protected]"},"sourceCommit":{"committedDate":"2024-09-26T13:18:28Z","message":"[EDR
Workflows] Add support to multiple agent id's in a policy for Osquery
(#193987)","sha":"2667db1a01db30e739ac09e8b1871b89df4e05a5","branchLabelMapping":{"^v9.0.0$":"main","^v8.16.0$":"8.x","^v(\\d+).(\\d+).\\d+$":"$1.$2"}},"sourcePullRequest":{"labels":["release_note:skip","v9.0.0","Team:Defend
Workflows","v8.16.0","backport:version"],"title":"[EDR Workflows] Add
support to multiple agent id's in a policy for
Osquery","number":193987,"url":"https://github.com/elastic/kibana/pull/193987","mergeCommit":{"message":"[EDR
Workflows] Add support to multiple agent id's in a policy for Osquery
(#193987)","sha":"2667db1a01db30e739ac09e8b1871b89df4e05a5"}},"sourceBranch":"main","suggestedTargetBranches":["8.x"],"targetPullRequestStates":[{"branch":"main","label":"v9.0.0","branchLabelMappingKey":"^v9.0.0$","isSourceBranch":true,"state":"MERGED","url":"https://github.com/elastic/kibana/pull/193987","number":193987,"mergeCommit":{"message":"[EDR
Workflows] Add support to multiple agent id's in a policy for Osquery
(#193987)","sha":"2667db1a01db30e739ac09e8b1871b89df4e05a5"}},{"branch":"8.x","label":"v8.16.0","branchLabelMappingKey":"^v8.16.0$","isSourceBranch":false,"state":"NOT_CREATED"}]}]
BACKPORT-->

Co-authored-by: Tomasz Ciecierski <[email protected]>

x-pack/plugins/osquery/public/agents/use_osquery_policies.ts

@@ -19,7 +19,7 @@ export const useOsqueryPolicies = () => {
   return useQuery(
     ['osqueryPolicies'],
     () =>
-      http.get<{ items: Array<{ policy_id: string; policy_ids: string[] }> }>(
+      http.get<{ items: Array<{ policy_ids: string[] }> }>(
         '/internal/osquery/fleet_wrapper/package_policies',
         { version: API_VERSIONS.internal.v1 }
       ),

x-pack/plugins/osquery/server/lib/fleet_integration.ts

@@ -21,13 +21,16 @@ export const getPackagePolicyDeleteCallback =
     ]);
     await Promise.all(
       map(deletedOsqueryManagerPolicies, async (deletedOsqueryManagerPolicy) => {
-        if (deletedOsqueryManagerPolicy.policy_id) {
+        const policyIds = deletedOsqueryManagerPolicy.policy_ids?.length
+          ? deletedOsqueryManagerPolicy.policy_ids
+          : ([deletedOsqueryManagerPolicy.policy_id] as string[]);
+        if (policyIds[0] !== undefined) {
           const foundPacks = await packsClient.find({
             type: packSavedObjectType,
-            hasReference: {
+            hasReference: policyIds.map((policyId: string) => ({
               type: LEGACY_AGENT_POLICY_SAVED_OBJECT_TYPE,
-              id: deletedOsqueryManagerPolicy.policy_id,
-            },
+              id: policyId,
+            })),
             perPage: 1000,
           });
 
@@ -43,15 +46,13 @@ export const getPackagePolicyDeleteCallback =
                   packSavedObjectType,
                   pack.id,
                   {
-                    shards: filter(
-                      pack.attributes.shards,
-                      (shard) => shard.key !== deletedOsqueryManagerPolicy.policy_id
+                    shards: filter(pack.attributes.shards, (shard) =>
+                      policyIds.includes(shard.key)
                     ),
                   },
                   {
-                    references: filter(
-                      pack.references,
-                      (reference) => reference.id !== deletedOsqueryManagerPolicy.policy_id
+                    references: filter(pack.references, (reference) =>
+                      policyIds.includes(reference.id)
                     ),
                   }
                 )

x-pack/plugins/osquery/server/routes/pack/create_pack_route.ts

@@ -152,9 +152,7 @@ export const createPackRoute = (router: IRouter, osqueryContext: OsqueryAppConte
                     }
 
                     set(draft, `inputs[0].config.osquery.value.packs.${packSO.attributes.name}`, {
-                      shard: policyShards[packagePolicy.policy_ids[0]] // TODO
-                        ? policyShards[packagePolicy.policy_ids[0]]
-                        : 100,
+                      shard: policyShards[agentPolicyId] ?? 100,
                       queries: convertSOQueriesToPackConfig(queries),
                     });
 

x-pack/plugins/osquery/server/routes/pack/update_pack_route.ts

@@ -7,7 +7,7 @@
 
 import moment from 'moment-timezone';
 import { set } from '@kbn/safer-lodash-set';
-import { unset, has, difference, filter, find, map, mapKeys, uniq, some, isEmpty } from 'lodash';
+import { unset, has, difference, filter, map, mapKeys, uniq, some, isEmpty } from 'lodash';
 import { produce } from 'immer';
 import type { PackagePolicy } from '@kbn/fleet-plugin/common';
 import {
@@ -194,8 +194,9 @@ export const updatePackRoute = (router: IRouter, osqueryContext: OsqueryAppConte
 
             await Promise.all(
               policyIds.map((agentPolicyId) => {
-                const packagePolicy = find(packagePolicies, ['policy_id', agentPolicyId]);
-
+                const packagePolicy = packagePolicies.find((policy) =>
+                  policy.policy_ids.includes(agentPolicyId)
+                );
                 if (packagePolicy) {
                   return packagePolicyService?.update(
                     internalSavedObjectsClient,
@@ -224,7 +225,9 @@ export const updatePackRoute = (router: IRouter, osqueryContext: OsqueryAppConte
           } else {
             await Promise.all(
               currentAgentPolicyIds.map((agentPolicyId) => {
-                const packagePolicy = find(currentPackagePolicies, ['policy_id', agentPolicyId]);
+                const packagePolicy = currentPackagePolicies.find((policy) =>
+                  policy.policy_ids.includes(agentPolicyId)
+                );
                 if (!packagePolicy) return;
 
                 return packagePolicyService?.update(
@@ -254,7 +257,9 @@ export const updatePackRoute = (router: IRouter, osqueryContext: OsqueryAppConte
 
           await Promise.all(
             agentPolicyIdsToRemove.map((agentPolicyId) => {
-              const packagePolicy = find(currentPackagePolicies, ['policy_id', agentPolicyId]);
+              const packagePolicy = currentPackagePolicies.find((policy) =>
+                policy.policy_ids.includes(agentPolicyId)
+              );
               if (packagePolicy) {
                 return packagePolicyService?.update(
                   internalSavedObjectsClient,
@@ -276,8 +281,9 @@ export const updatePackRoute = (router: IRouter, osqueryContext: OsqueryAppConte
 
           await Promise.all(
             agentPolicyIdsToUpdate.map((agentPolicyId) => {
-              const packagePolicy = find(packagePolicies, ['policy_id', agentPolicyId]);
-
+              const packagePolicy = packagePolicies.find((policy) =>
+                policy.policy_ids.includes(agentPolicyId)
+              );
               if (packagePolicy) {
                 return packagePolicyService?.update(
                   internalSavedObjectsClient,
@@ -296,9 +302,7 @@ export const updatePackRoute = (router: IRouter, osqueryContext: OsqueryAppConte
                       draft,
                       `inputs[0].config.osquery.value.packs.${updatedPackSO.attributes.name}`,
                       {
-                        shard: policyShards[packagePolicy.policy_ids[0]] // TODO
-                          ? policyShards[packagePolicy.policy_ids[0]]
-                          : 100,
+                        shard: policyShards[agentPolicyId] ?? 100,
                         queries: convertSOQueriesToPackConfig(updatedPackSO.attributes.queries),
                       }
                     );
@@ -331,9 +335,7 @@ export const updatePackRoute = (router: IRouter, osqueryContext: OsqueryAppConte
                       draft,
                       `inputs[0].config.osquery.value.packs.${updatedPackSO.attributes.name}`,
                       {
-                        shard: policyShards[packagePolicy.policy_ids[0]] // TODO
-                          ? policyShards[packagePolicy.policy_ids[0]]
-                          : 100,
+                        shard: policyShards[agentPolicyId] ?? 100,
                         queries: convertSOQueriesToPackConfig(updatedPackSO.attributes.queries),
                       }
                     );