Skip to content

Commit 55bca3e

Browse files
MarfuenMarfuen
authored
Merge pull request #874 from trycompai/mariano/comp-166-osquery-agent-integration
[dev] [Marfuen] mariano/comp-166-osquery-agent-integration
2 parents b577aa9 + 60870c4 commit 55bca3e

9 files changed

Lines changed: 491 additions & 235 deletions

File tree

‎apps/app/src/components/mobile-menu.tsx‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,10 @@ export function MobileMenu({
3939
metadata: org.metadata ? String(org.metadata) : null,
4040
stripeCustomerId: null,
4141
website: null,
42+
fleetDmSecret: "",
43+
fleetDmLabelId: null,
44+
osqueryAgentDownloadUrl: null,
45+
isFleetSetupCompleted: false,
4246
}));
4347

4448
const currentOrganization =

‎apps/app/src/jobs/tasks/device/generate-agent-file.ts‎

Lines changed: 155 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,155 @@
1+
import { logger, task } from "@trigger.dev/sdk/v3";
2+
import { db } from "@comp/db";
3+
import { promisify } from "node:util";
4+
import { exec as callbackExec } from "node:child_process";
5+
import { createReadStream, existsSync, mkdtempSync, rmSync } from "node:fs";
6+
import { tmpdir } from "node:os";
7+
import path from "node:path";
8+
import { PutObjectCommand } from "@aws-sdk/client-s3";
9+
import { s3Client } from "@/app/s3";
10+
import { fleet } from "@/lib/fleet";
11+
12+
export const generateAgentFile = task({
13+
id: "generate-agent-file",
14+
retry: {
15+
maxAttempts: 3,
16+
},
17+
cleanup: async ({ organizationId }: { organizationId: string }) => {
18+
// Delete the tmp dir.
19+
const tmpDir = path.join(tmpdir(), `pkg-${organizationId}-`);
20+
if (existsSync(tmpDir)) {
21+
rmSync(tmpDir, { recursive: true });
22+
}
23+
},
24+
run: async ({ organizationId }: { organizationId: string }) => {
25+
const organization = await db.organization.findUnique({
26+
where: {
27+
id: organizationId,
28+
},
29+
});
30+
31+
if (!organization) {
32+
logger.error(`Organization ${organizationId} not found`);
33+
return;
34+
}
35+
36+
if (organization.isFleetSetupCompleted) {
37+
logger.info(`Organization ${organizationId} already has fleet set up`);
38+
return;
39+
}
40+
41+
const fleetDevicePathMac = process.env.FLEET_DEVICE_PATH_MAC;
42+
if (!fleetDevicePathMac) {
43+
logger.error("FLEET_DEVICE_PATH_MAC not configured");
44+
return;
45+
}
46+
47+
// Create a manual label that we can assign to hosts.
48+
const response = await fleet.post("/labels", {
49+
name: organization.id,
50+
query: `SELECT 1 FROM file WHERE path = '${fleetDevicePathMac}/${organizationId}' LIMIT 1;`,
51+
});
52+
53+
const enrollSecret = organization.fleetDmSecret;
54+
55+
const secretsResponse = await fleet.get("/spec/enroll_secret");
56+
const existingSecrets = secretsResponse.data.spec.secrets;
57+
58+
await fleet.post("/spec/enroll_secret", {
59+
spec: {
60+
secrets: [
61+
...existingSecrets,
62+
{
63+
secret: enrollSecret,
64+
name: organization.id,
65+
},
66+
],
67+
},
68+
});
69+
70+
if (!response.data) {
71+
logger.error(`Failed to create label for organization ${organizationId}`);
72+
return;
73+
}
74+
75+
// Store label ID in organization.
76+
await db.organization.update({
77+
where: {
78+
id: organizationId,
79+
},
80+
data: {
81+
fleetDmLabelId: response.data.label.id,
82+
},
83+
});
84+
85+
// Create osquery agent file.
86+
const execAsync = promisify(callbackExec);
87+
const fleetUrl = process.env.FLEET_URL;
88+
89+
if (!enrollSecret) {
90+
logger.error(
91+
"FLEET_ENROLL_SECRET is not set. Cannot create osquery agent."
92+
);
93+
return;
94+
}
95+
96+
try {
97+
const workDir = mkdtempSync(
98+
path.join(tmpdir(), `pkg-${organizationId}-`)
99+
);
100+
101+
logger.info(`Building .pkg in ${workDir}`);
102+
103+
const commandMac = `fleetctl package \
104+
--type=pkg \
105+
--fleet-url ${fleetUrl} \
106+
--enable-scripts \
107+
--fleet-desktop \
108+
--verbose \
109+
--enroll-secret "${enrollSecret}"`;
110+
111+
logger.info(`Executing; command: ${commandMac}`);
112+
113+
await execAsync(commandMac, {
114+
cwd: workDir,
115+
});
116+
117+
const pkgPath = path.join(workDir, "fleet-osquery.pkg");
118+
119+
logger.info(`Created fleet-osquery.pkg in ${pkgPath}`);
120+
121+
const s3KeyPkg = `${organizationId}/macos/fleet-osquery.pkg`;
122+
123+
// Upload the zip to S3
124+
const putObjectCommandPkg = new PutObjectCommand({
125+
Bucket: "compai-fleet-packages",
126+
Key: s3KeyPkg,
127+
Body: createReadStream(pkgPath),
128+
ContentType: "application/octet-stream",
129+
});
130+
131+
logger.info(`Uploading fleet-osquery.pkg to S3: ${s3KeyPkg}`);
132+
await s3Client.send(putObjectCommandPkg);
133+
134+
const s3Region = await s3Client.config.region();
135+
const s3ObjectUrlPkg = `https://compai-fleet-packages.s3.${s3Region}.amazonaws.com/${s3KeyPkg}`;
136+
137+
logger.info("S3 Upload successful.", {
138+
fileUrlPkg: s3ObjectUrlPkg,
139+
});
140+
141+
await db.organization.update({
142+
where: { id: organizationId },
143+
data: {
144+
osqueryAgentDownloadUrl: s3ObjectUrlPkg,
145+
isFleetSetupCompleted: true,
146+
},
147+
});
148+
logger.info(`Stored S3 bundle URL for organization ${organizationId}`);
149+
} catch (error) {
150+
logger.error("Error in fleetctl packaging or S3 upload process", {
151+
error,
152+
});
153+
}
154+
},
155+
});

‎apps/app/src/jobs/tasks/device/generate-agent-for-all-orgs.ts‎

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,27 @@
1+
import { db } from "@comp/db";
2+
import { logger, task } from "@trigger.dev/sdk/v3";
3+
import { generateAgentFile } from "./generate-agent-file";
4+
5+
export const generateAgentForAllOrgs = task({
6+
id: "generate-agent-for-all-orgs",
7+
run: async () => {
8+
const organizations = await db.organization.findMany({
9+
where: {
10+
isFleetSetupCompleted: false,
11+
},
12+
});
13+
14+
logger.info(
15+
`Found ${organizations.length} organizations to generate agent for`
16+
);
17+
18+
const batchItems = organizations.map((organization) => ({
19+
payload: {
20+
organizationId: organization.id,
21+
},
22+
}));
23+
24+
logger.info(`Triggering batch job for ${batchItems.length} organizations`);
25+
await generateAgentFile.batchTrigger(batchItems);
26+
},
27+
});

‎apps/app/src/jobs/tasks/device/generate-agent-for-orgs-skipped-onboarding.ts‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,32 @@
1+
import { db } from "@comp/db";
2+
import { logger, schedules } from "@trigger.dev/sdk/v3";
3+
import { generateAgentFile } from "./generate-agent-file";
4+
5+
export const generateAgentForOrgsSkippedOnboarding = schedules.task({
6+
id: "generate-agent-for-orgs-skipped-onboarding",
7+
cron: "0 * * * *", // every hour
8+
maxDuration: 1000 * 60 * 10, // 10 minutes
9+
run: async () => {
10+
const organizations = await db.organization.findMany({
11+
where: {
12+
isFleetSetupCompleted: false,
13+
onboarding: {
14+
completed: true,
15+
},
16+
},
17+
});
18+
19+
logger.info(
20+
`Found ${organizations.length} organizations to generate agent for`
21+
);
22+
23+
const batchItems = organizations.map((organization) => ({
24+
payload: {
25+
organizationId: organization.id,
26+
},
27+
}));
28+
29+
logger.info(`Triggering batch job for ${batchItems.length} organizations`);
30+
await generateAgentFile.batchTrigger(batchItems);
31+
},
32+
});

0 commit comments

Comments
 (0)