From e4348f2ec28a1b58cb43db987454f4236af960e6 Mon Sep 17 00:00:00 2001
From: Michael Wedl <michael@syslifters.com>
Date: Wed, 17 Apr 2024 07:47:48 +0200
Subject: [PATCH] Fix user notes

---
 frontend/src/stores/project.ts   | 16 +++++++++-------
 frontend/src/stores/usernotes.ts |  3 +--
 frontend/src/utils/collab.ts     |  6 +++---
 3 files changed, 13 insertions(+), 12 deletions(-)

diff --git a/frontend/src/stores/project.ts b/frontend/src/stores/project.ts
index 8959accab..7f393cf3d 100644
--- a/frontend/src/stores/project.ts
+++ b/frontend/src/stores/project.ts
@@ -51,7 +51,11 @@ export const useProjectStore = defineStore('project', {
         // notes: ProjectNote[],
         getByIdSync: Promise<PentestProject> | null,
         notesCollabState: CollabStoreState<{ notes: {[key: string]: ProjectNote}}>,
-        reportingCollabState: CollabStoreState<{ findings: {[key: string]: PentestFinding }, sections: {[key: string]: ReportSection } }>,
+        reportingCollabState: CollabStoreState<{ 
+          project: {id: string, project_type: string, override_finding_order: boolean}, 
+          findings: {[key: string]: PentestFinding }, 
+          sections: {[key: string]: ReportSection } 
+        }>,
       }
     },
   }),
@@ -102,9 +106,8 @@ export const useProjectStore = defineStore('project', {
           getByIdSync: null,
           notesCollabState: makeCollabStoreState({
             websocketPath: `/ws/pentestprojects/${projectId}/notes/`,
-            initialData: { notes: {} },
-            handleAdditionalWebSocketMessages: (msgData: any) => {
-              const collabState = this.data[projectId].notesCollabState;
+            initialData: { notes: {} as {[key: string]: ProjectNote} },
+            handleAdditionalWebSocketMessages: (msgData: any, collabState) => {
               if (msgData.type === CollabEventType.SORT && msgData.path === 'notes') {
                 for (const note of Object.values(collabState.data.notes)) {
                   const no = msgData.sort.find((n: ProjectNote) => n.id === note.id);
@@ -119,9 +122,8 @@ export const useProjectStore = defineStore('project', {
           }),
           reportingCollabState: makeCollabStoreState({
             websocketPath: `/ws/pentestprojects/${projectId}/reporting/`,
-            initialData: { project: {}, findings: {}, sections: {} },
-            handleAdditionalWebSocketMessages: (msgData: any) => {
-              const collabState = this.data[projectId].reportingCollabState;
+            initialData: { project: {} as any, findings: {} as {[key: string]: PentestFinding}, sections: {} as {[key: string]: ReportSection} },
+            handleAdditionalWebSocketMessages: (msgData: any, collabState) => {
               if (msgData.type === CollabEventType.SORT && msgData.path === 'findings') {
                 for (const finding of Object.values(collabState.data.findings)) {
                   const fo = msgData.sort.find((n: PentestFinding) => n.id === finding.id);
diff --git a/frontend/src/stores/usernotes.ts b/frontend/src/stores/usernotes.ts
index 3567901ec..8c01db378 100644
--- a/frontend/src/stores/usernotes.ts
+++ b/frontend/src/stores/usernotes.ts
@@ -39,8 +39,7 @@ export const useUserNotesStore = defineStore('usernotes', {
       notesCollabState: makeCollabStoreState({
         websocketPath: '/ws/pentestusers/self/notes/',
         initialData: { notes: {} as {[key: string]: UserNote} },
-        handleAdditionalWebSocketMessages: (msgData: any) => {
-          const collabState = (this as any).notesCollabState;
+        handleAdditionalWebSocketMessages: (msgData: any, collabState) => {
           if (msgData.type === CollabEventType.SORT && msgData.path === 'notes') {
             for (const note of Object.values(collabState.data.notes) as UserNote[]) {
               const no = msgData.sort.find((n: UserNote) => n.id === note.id);
diff --git a/frontend/src/utils/collab.ts b/frontend/src/utils/collab.ts
index fb3b12c59..3178b37eb 100644
--- a/frontend/src/utils/collab.ts
+++ b/frontend/src/utils/collab.ts
@@ -43,7 +43,7 @@ export type CollabStoreState<T> = {
   websocketPath: string;
   connectionError?: { error: any, message?: string };
   websocketConnectionLostTimeout?: ReturnType<typeof throttle>;
-  handleAdditionalWebSocketMessages?: (event: any) => boolean;
+  handleAdditionalWebSocketMessages?: (event: any, collabState: CollabStoreState<T>) => boolean;
   perPathState: Map<string, {
     sendUpdateTextThrottled: ReturnType<typeof throttle>;
     unconfirmedTextUpdates: TextUpdate[];
@@ -74,7 +74,7 @@ export type CollabStoreState<T> = {
 export function makeCollabStoreState<T>(options: {
   websocketPath: string, 
   initialData: T,
-  handleAdditionalWebSocketMessages?: (event: any) => boolean
+  handleAdditionalWebSocketMessages?: (event: any, storeState: CollabStoreState<T>) => boolean
 }): CollabStoreState<T> {
   return {
     data: options.initialData,
@@ -166,7 +166,7 @@ export function useCollab<T = any>(storeState: CollabStoreState<T>) {
           storeState.version = msgData.version;
         }
 
-        if (storeState.handleAdditionalWebSocketMessages?.(msgData)) {
+        if (storeState.handleAdditionalWebSocketMessages?.(msgData, storeState)) {
           // Already handled
         } else if (msgData.type === CollabEventType.INIT) {
           storeState.connectionState = CollabConnectionState.OPEN;