From 51e1fffa4800fd76baade9c98bc15f00a78ff2de Mon Sep 17 00:00:00 2001
From: Austin Golding <austin.golding.dev@icloud.com>
Date: Tue, 22 Aug 2023 18:07:55 -0700
Subject: [PATCH] Add ability to view cobalt strike raw log files (#176)

---
 .../client/src/store/campaign/campaign.ts     | 11 ++++
 .../src/store/graphql/RootStore.base.ts       | 15 +++++
 .../client/src/store/util/cleaned-queries.ts  |  1 +
 .../Explore/Panels/Command/RawLogsDialog.tsx  | 66 ++++++++++++++-----
 applications/server/schema.graphql            |  3 +
 applications/server/src/config.ts             |  4 +-
 .../server/src/store/command-resolvers.ts     |  2 +-
 .../server/src/store/log-resolvers.ts         | 23 +++++++
 .../server/src/store/operator-resolvers.ts    |  4 +-
 9 files changed, 110 insertions(+), 19 deletions(-)

diff --git a/applications/client/src/store/campaign/campaign.ts b/applications/client/src/store/campaign/campaign.ts
index dd994324..77cb5e14 100644
--- a/applications/client/src/store/campaign/campaign.ts
+++ b/applications/client/src/store/campaign/campaign.ts
@@ -244,4 +244,15 @@ export class CampaignStore extends ExtendedModel(() => ({
 			direction: direction as SortDirection,
 		};
 	}
+
+	@computed get parsers(): string[] {
+		if (this.id)
+			return (
+				(this.appStore?.graphqlStore.campaigns
+					.get(this.id)
+					?.parsers?.map((parser) => parser.parserName)
+					.filter(Boolean) as string[]) || []
+			);
+		return [];
+	}
 }
diff --git a/applications/client/src/store/graphql/RootStore.base.ts b/applications/client/src/store/graphql/RootStore.base.ts
index 88b4f4a8..e5e09cb8 100644
--- a/applications/client/src/store/graphql/RootStore.base.ts
+++ b/applications/client/src/store/graphql/RootStore.base.ts
@@ -176,6 +176,7 @@ export enum RootStoreBaseQueries {
 	queryHosts = 'queryHosts',
 	queryImages = 'queryImages',
 	queryLinks = 'queryLinks',
+	queryLogFilesByBeaconId = 'queryLogFilesByBeaconId',
 	queryLogs = 'queryLogs',
 	queryLogsByBeaconId = 'queryLogsByBeaconId',
 	queryNonHidableEntities = 'queryNonHidableEntities',
@@ -633,6 +634,20 @@ export class RootStoreBase extends ExtendedModel(
 			!!clean
 		);
 	}
+	// Get logs from beacon sorted by time. The goal is to be able to re-create the full log for that beacon.
+	@modelAction queryLogFilesByBeaconId(
+		variables: { beaconId: string; campaignId: string },
+		_?: any,
+		options: QueryOptions = {},
+		clean?: boolean
+	) {
+		return this.query<{ logFilesByBeaconId: StringModel[] }>(
+			`query logFilesByBeaconId($beaconId: String!, $campaignId: String!) { logFilesByBeaconId(beaconId: $beaconId, campaignId: $campaignId)  }`,
+			variables,
+			options,
+			!!clean
+		);
+	}
 	// Get log entries by ids
 	@modelAction queryLogs(
 		variables: { beaconId?: string; campaignId: string; hostId?: string },
diff --git a/applications/client/src/store/util/cleaned-queries.ts b/applications/client/src/store/util/cleaned-queries.ts
index 0309e40e..77e13157 100644
--- a/applications/client/src/store/util/cleaned-queries.ts
+++ b/applications/client/src/store/util/cleaned-queries.ts
@@ -70,6 +70,7 @@ export const hostQuery = hostModelPrimitives.cobaltStrikeServer.beaconIds
 export const campaignsQuery = campaignModelPrimitives
 	.lastOpenedBy((user) => user)
 	.creator((user) => user)
+	.parsers((parser) => parser.parserName)
 	.toString();
 
 export const rawLogOutputQuery = logEntryModelPrimitives;
diff --git a/applications/client/src/views/Campaign/Explore/Panels/Command/RawLogsDialog.tsx b/applications/client/src/views/Campaign/Explore/Panels/Command/RawLogsDialog.tsx
index 8a461f4f..ca73a72c 100644
--- a/applications/client/src/views/Campaign/Explore/Panels/Command/RawLogsDialog.tsx
+++ b/applications/client/src/views/Campaign/Explore/Panels/Command/RawLogsDialog.tsx
@@ -1,6 +1,5 @@
-import type {} from '@blueprintjs/core';
 import { Button, ButtonGroup, Divider, NonIdealState, Spinner } from '@blueprintjs/core';
-import { Copy16, UpToTop16 } from '@carbon/icons-react';
+import { Copy16, UpToTop16, Document16 } from '@carbon/icons-react';
 import { css } from '@emotion/react';
 import type { DialogExProps } from '@redeye/client/components';
 import { CarbonIcon, DialogEx } from '@redeye/client/components';
@@ -26,6 +25,7 @@ export const RawLogsDialog = observer<RawLogsViewerProps>(({ ...props }) => {
 	const scrollTopTargetRef = useRef<HTMLDivElement>(null);
 
 	const state = createState({
+		showLogFile: false,
 		get isOpen(): boolean {
 			return (
 				!!(this.beacon && store.router.queryParams['raw-logs']) ||
@@ -46,6 +46,7 @@ export const RawLogsDialog = observer<RawLogsViewerProps>(({ ...props }) => {
 			return this.command?.input?.current?.id;
 		},
 		onClose() {
+			this.showLogFile = false;
 			store.router.updateQueryParams({ queryParams: { 'raw-logs': undefined, 'raw-command': undefined } });
 		},
 		handleCopyText(logsByBeaconId) {
@@ -79,6 +80,23 @@ export const RawLogsDialog = observer<RawLogsViewerProps>(({ ...props }) => {
 		{ enabled: state.isOpen }
 	);
 
+	const {
+		data: { logFilesByBeaconId } = {},
+		isLoading: isLoadingFiles,
+		isError: isErrorFiles,
+	} = useQuery(
+		['rawLogFile', store.campaign.id, state.commandInputId],
+		async () =>
+			await store.graphqlStore.queryLogFilesByBeaconId(
+				{
+					campaignId: store.campaign.id!,
+					beaconId: state.beacon?.id!,
+				},
+				selectFromLogEntry().toString()
+			),
+		{ enabled: state.isOpen && state.showLogFile && store.campaign.parsers.includes('cobalt-strike-parser') }
+	);
+
 	if (!state.command && !state.beacon) return null;
 
 	return (
@@ -106,6 +124,16 @@ export const RawLogsDialog = observer<RawLogsViewerProps>(({ ...props }) => {
 							</Txt>
 						</PanelHeader>
 						<ButtonGroup>
+							{store.campaign.parsers.includes('cobalt-strike-parser') ? (
+								<Button
+									cy-test="view-log-file"
+									text="View Log File"
+									rightIcon={<CarbonIcon icon={Document16} />}
+									active={state.showLogFile}
+									onClick={() => state.update('showLogFile', !state.showLogFile)}
+									minimal
+								/>
+							) : null}
 							<Button
 								cy-test="scroll-to-top"
 								text="Scroll To Top"
@@ -127,21 +155,29 @@ export const RawLogsDialog = observer<RawLogsViewerProps>(({ ...props }) => {
 			}
 		>
 			<div ref={scrollTopTargetRef} />
-			{isLoading && <Spinner css={messagePaddingStyles} />}
-			{isError && <NonIdealState title="Unable to fetch Logs" icon="error" css={messagePaddingStyles} />}
+			{(isLoading || (state.showLogFile && isLoadingFiles)) && <Spinner css={messagePaddingStyles} />}
+			{(isError || (state.showLogFile && isErrorFiles)) && (
+				<NonIdealState title="Unable to fetch Logs" icon="error" css={messagePaddingStyles} />
+			)}
 			<div cy-test="log" css={outputScrollWrapperStyle}>
 				<div css={outputOverflowWrapperStyle}>
-					{logsByBeaconId
-						?.filter((logEntry) => !!getTextFromLog(logEntry))
-						.map((logEntry) => {
-							const isScrollTarget = state.commandLogEntryIds?.includes?.(logEntry.id);
-							const Component = isScrollTarget ? ScrollTargetPre : 'pre';
-							return (
-								<Component key={logEntry.id} cy-test="log-entry" css={[preStyles]}>
-									{getTextFromLog(logEntry)}
-								</Component>
-							);
-						})}
+					{state.showLogFile
+						? logFilesByBeaconId?.map((file) => (
+								<pre cy-test="log-entry" css={[preStyles]}>
+									{file}
+								</pre>
+						  ))
+						: logsByBeaconId
+								?.filter((logEntry) => !!getTextFromLog(logEntry))
+								.map((logEntry) => {
+									const isScrollTarget = state.commandLogEntryIds?.includes?.(logEntry.id);
+									const Component = isScrollTarget ? ScrollTargetPre : 'pre';
+									return (
+										<Component key={logEntry.id} cy-test="log-entry" css={[preStyles]}>
+											{getTextFromLog(logEntry)}
+										</Component>
+									);
+								})}
 				</div>
 			</div>
 		</DialogEx>
diff --git a/applications/server/schema.graphql b/applications/server/schema.graphql
index c7470026..9dc3d98b 100644
--- a/applications/server/schema.graphql
+++ b/applications/server/schema.graphql
@@ -534,6 +534,9 @@ type Query {
   """Get all links"""
   links(campaignId: String!, hidden: Boolean = false): [Link]
 
+  """Get log files from beacon"""
+  logFilesByBeaconId(beaconId: String!, campaignId: String!): [String!]!
+
   """Get log entries by ids"""
   logs(beaconId: String, campaignId: String!, hostId: String): [LogEntry]
 
diff --git a/applications/server/src/config.ts b/applications/server/src/config.ts
index ea7cb6f3..acadcf80 100644
--- a/applications/server/src/config.ts
+++ b/applications/server/src/config.ts
@@ -46,9 +46,9 @@ const castConfig = (cliArgs: cliArgs): ConfigDefinition => {
 		overrides.production = false;
 		overrides.databaseMode = DatabaseMode.DEV_PERSIST;
 	}
-	if (cliArgs.port) overrides.port = cliArgs.port;
+	if (cliArgs.port) overrides.port = parseInt(`${cliArgs.port}`, 10);
 	if (cliArgs.redTeam) overrides.redTeam = true;
-	if (cliArgs.childProcesses) overrides.maxParserSubprocesses = cliArgs.childProcesses;
+	if (cliArgs.childProcesses) overrides.maxParserSubprocesses = parseInt(`${cliArgs.childProcesses}`, 10);
 	if (cliArgs.password) overrides.password = cliArgs.password;
 	if (cliArgs.parsers) {
 		if (cliArgs.parsers === true) {
diff --git a/applications/server/src/store/command-resolvers.ts b/applications/server/src/store/command-resolvers.ts
index 32d24425..5dec5f7e 100644
--- a/applications/server/src/store/command-resolvers.ts
+++ b/applications/server/src/store/command-resolvers.ts
@@ -212,7 +212,7 @@ export class CommandTypeCountResolvers {
 		hidden: boolean = false
 	): Promise<CommandTypeCount[]> {
 		const em = await connectToProjectEmOrFail(campaignId, ctx);
-		const commands = await em.find(Command, beaconHidden(hidden), { populate: false });
+		const commands = await em.find(Command, beaconHidden(hidden), { populate: ['commandGroups'] });
 		const countObj = commands.reduce<Record<string, CountObjItem>>((acc, current) => {
 			if (acc[current.inputText]) {
 				acc[current.inputText] = {
diff --git a/applications/server/src/store/log-resolvers.ts b/applications/server/src/store/log-resolvers.ts
index 109c7d88..a7d1f21c 100644
--- a/applications/server/src/store/log-resolvers.ts
+++ b/applications/server/src/store/log-resolvers.ts
@@ -4,6 +4,7 @@ import { connectToProjectEmOrFail } from './utils/project-db';
 import { RelationPath } from './utils/relation-path';
 import type { Relation } from './utils/relation-path';
 import type { GraphQLContext } from '../types';
+import { readFileSync } from 'fs-extra';
 
 @Resolver(LogEntry)
 export class LogResolvers {
@@ -56,4 +57,26 @@ export class LogResolvers {
 
 		return logs;
 	}
+
+	@Authorized()
+	@Query(() => [String], {
+		description: 'Get log files from beacon',
+	})
+	async logFilesByBeaconId(
+		@Ctx() ctx: GraphQLContext,
+		@Arg('campaignId', () => String) campaignId: string,
+		@Arg('beaconId', () => String) beaconId: string
+	): Promise<string[]> {
+		const em = await connectToProjectEmOrFail(campaignId, ctx);
+		const logs: LogEntry[] = await em.createQueryBuilder(LogEntry).where({ beacon: beaconId });
+		const files: string[] = [];
+		const seenFiles = new Set<string>();
+		for (const log of logs) {
+			if (log.filepath && !seenFiles.has(log.filepath)) {
+				files.push(readFileSync(log.filepath, 'utf8'));
+				seenFiles.add(log.filepath);
+			}
+		}
+		return files;
+	}
 }
diff --git a/applications/server/src/store/operator-resolvers.ts b/applications/server/src/store/operator-resolvers.ts
index 233bcbae..19f4b634 100644
--- a/applications/server/src/store/operator-resolvers.ts
+++ b/applications/server/src/store/operator-resolvers.ts
@@ -18,7 +18,9 @@ export class OperatorResolvers {
 		@RelationPath() relationPaths: Relation<Operator>
 	): Promise<Operator[]> {
 		const em = await connectToProjectEmOrFail(campaignId, ctx);
-		const operators = await em.find(Operator, !hidden ? { beacons: { hidden } } : {}, { populate: relationPaths });
+		const operators = await em.find(Operator, !hidden ? { beacons: { hidden } } : {}, {
+			populate: [...relationPaths, 'beacons'],
+		});
 		for (const operator of operators) {
 			await operator.beacons.init({ populate: false, where: defaultHidden(hidden) });
 			await operator.commands.init({ populate: false, where: beaconHidden(hidden) });