Implemented Bookmarks Feature

cmd/report_list.go

@@ -20,8 +20,9 @@ import (
 )
 
 var listCmdFlags = struct {
-	DbURI    string
-	JsonFile string
+	DbURI     string
+	JsonFile  string
+	Bookmarks bool
 }{}
 var listCmd = &cobra.Command{
 	Use:   "list",
@@ -87,7 +88,13 @@ lines file.`)),
 			return
 		}
 
-		if err := conn.Model(&models.Result{}).Preload(clause.Associations).Find(&results).Error; err != nil {
+		if listCmdFlags.Bookmarks {
+			err = conn.Model(&models.Result{}).Preload(clause.Associations).Where("bookmarked = true").Find(&results).Error
+		} else {
+			err = conn.Model(&models.Result{}).Preload(clause.Associations).Find(&results).Error
+		}
+
+		if err != nil {
 			log.Error("could not get list", "err", err)
 			return
 		}
@@ -101,6 +108,7 @@ func init() {
 
 	listCmd.Flags().StringVar(&listCmdFlags.DbURI, "db-uri", "sqlite://gowitness.sqlite3", "The location of a gowitness database")
 	listCmd.Flags().StringVar(&listCmdFlags.JsonFile, "json-file", "", "The location of a JSON Lines results file (e.g., ./gowitness.jsonl). This flag takes precedence over --db-uri")
+	listCmd.Flags().BoolVar(&listCmdFlags.Bookmarks, "bookmarks", false, "Only list bookmarked results")
 }
 
 func renderTable(results []*models.Result) {

pkg/models/models.go

@@ -28,6 +28,7 @@ type Result struct {
 	PerceptionHash        string    `json:"perception_hash" gorm:"index"`
 	PerceptionHashGroupId uint      `json:"perception_hash_group_id" gorm:"index"`
 	Screenshot            string    `json:"screenshot"`
+	Bookmarked            bool      `json:"bookmarked"`
 
 	// Name of the screenshot file
 	Filename string `json:"file_name"`

web/api/bookmark.go

@@ -0,0 +1,61 @@
+package api
+
+import (
+	"encoding/json"
+	"net/http"
+
+	"github.com/sensepost/gowitness/pkg/log"
+	"github.com/sensepost/gowitness/pkg/models"
+)
+
+type bookmarkRequest struct {
+	ID int `json:"id"`
+}
+
+// BookmarkHandler inverts the state of a bookmark
+//
+//	@Summary		Bookmark/Unbookmarks result
+//	@Description	Inverts the bookmark status of a result, writing results to the database.
+//	@Tags			Results
+//	@Accept			json
+//	@Produce		json
+//	@Param			query	body	bookmarkRequest	true	"The bookmark request object"
+//	@Success		200		{string}	string			"bookmarked"
+//	@Router			/results/bookmark [post]
+func (h *ApiHandler) BookmarkHandler(w http.ResponseWriter, r *http.Request) {
+	var request bookmarkRequest
+	if err := json.NewDecoder(r.Body).Decode(&request); err != nil {
+		log.Error("failed to read json request", "err", err)
+		http.Error(w, "Error reading JSON request", http.StatusInternalServerError)
+		return
+	}
+
+	var bookmarked bool
+	if err := h.DB.Model(&models.Result{}).Where("id = ?", request.ID).Select("bookmarked").First(&bookmarked).Error; err != nil {
+		log.Error("failed to get bookmark status", "err", err)
+		http.Error(w, "Error getting result bookmark value", http.StatusInternalServerError)
+		return
+	}
+
+	log.Info("inverting bookmark id", "id", request.ID)
+	if err := h.DB.Model(&models.Result{}).Where("id = ?", request.ID).Update("bookmarked", !bookmarked).Error; err != nil {
+		log.Error("failed to update result bookmark", "err", err)
+		http.Error(w, "Error updating result bookmark value", http.StatusInternalServerError)
+		return
+	}
+
+	var response string
+	if bookmarked {
+		response = `removed bookmark`
+	} else {
+		response = `bookmarked`
+	}
+
+	jsonData, err := json.Marshal(response)
+	if err != nil {
+		http.Error(w, "Error creating JSON response", http.StatusInternalServerError)
+		return
+	}
+
+	w.Write(jsonData)
+}

web/api/gallery.go

@@ -28,6 +28,7 @@ type galleryContent struct {
 	Screenshot   string    `json:"screenshot"`
 	Failed       bool      `json:"failed"`
 	Technologies []string  `json:"technologies"`
+	Bookmarked   bool      `json:"bookmarked"`
 }
 
 // GalleryHandler gets a paginated gallery
@@ -43,6 +44,7 @@ type galleryContent struct {
 //	@Param			status			query		string	false	"A comma seperated list of HTTP status codes to filter by."
 //	@Param			perception		query		boolean	false	"Order the results by perception hash."
 //	@Param			failed			query		boolean	false	"Include failed screenshots in the results."
+//	@Param			bookmarked		query		boolean	false	"Only return bookmarked items"
 //	@Success		200				{object}	galleryResponse
 //	@Router			/results/gallery [get]
 func (h *ApiHandler) GalleryHandler(w http.ResponseWriter, r *http.Request) {
@@ -98,6 +100,13 @@ func (h *ApiHandler) GalleryHandler(w http.ResponseWriter, r *http.Request) {
 		showFailed = true
 	}
 
+	// bookmarked filtering
+	var bookmarked bool
+	bookmarked, err = strconv.ParseBool(r.URL.Query().Get("bookmarked"))
+	if err != nil {
+		bookmarked = false
+	}
+
 	// query the db
 	var queryResults []*models.Result
 	query := h.DB.Model(&models.Result{}).Limit(results.Limit).
@@ -121,6 +130,10 @@ func (h *ApiHandler) GalleryHandler(w http.ResponseWriter, r *http.Request) {
 		query.Where("failed = ?", showFailed)
 	}
 
+	if bookmarked {
+		query.Where("bookmarked = ?", bookmarked)
+	}
+
 	// run the query
 	if err := query.Find(&queryResults).Error; err != nil {
 		log.Error("could not get gallery", "err", err)
@@ -145,6 +158,7 @@ func (h *ApiHandler) GalleryHandler(w http.ResponseWriter, r *http.Request) {
 			Screenshot:   result.Screenshot,
 			Failed:       result.Failed,
 			Technologies: technologies,
+			Bookmarked:   result.Bookmarked,
 		})
 	}
 

web/api/list.go

@@ -18,6 +18,7 @@ type listResponse struct {
 	Protocol       string `json:"protocol"`
 	ContentLength  int64  `json:"content_length"`
 	Title          string `json:"title"`
+	Bookmarked     bool   `json:"bookmarked"`
 
 	// Failed flag set if the result should be considered failed
 	Failed       bool   `json:"failed"`

web/api/search.go

@@ -29,12 +29,13 @@ type searchResult struct {
 	FailedReason   string   `json:"failed_reason"`
 	Filename       string   `json:"file_name"`
 	Screenshot     string   `json:"screenshot"`
+	Bookmarked     bool     `json:"bookmarked"`
 	MatchedFields  []string `json:"matched_fields"`
 }
 
 // searchOperators are the operators we support. everything else is
 // "free text"
-var searchOperators = []string{"title", "body", "tech", "header", "p"}
+var searchOperators = []string{"title", "body", "tech", "header", "p", "bookmarked"}
 
 // SearchHandler handles search
 //

web/server.go

@@ -79,6 +79,7 @@ func (s *Server) Run() {
 		r.Get("/results/gallery", apih.GalleryHandler)
 		r.Get("/results/list", apih.ListHandler)
 		r.Get("/results/detail/{id}", apih.DetailHandler)
+		r.Post("/results/bookmark", apih.BookmarkHandler)
 		r.Post("/results/delete", apih.DeleteResultHandler)
 		r.Get("/results/technology", apih.TechnologyListHandler)
 	})

web/ui/src/lib/api/api.ts

@@ -47,6 +47,10 @@ const endpoints = {
     path: `/search`,
     returnas: {} as searchresult
   },
+  bookmark: {
+    path: `/results/bookmark`,
+    returnas: "" as string
+  },
   delete: {
     path: `/results/delete`,
     returnas: "" as string

web/ui/src/lib/api/bookmark.ts

@@ -0,0 +1,23 @@
+import { toast } from "@/hooks/use-toast";
+import * as api from "@/lib/api/api";
+
+const bookmarkResult = async (id: number): Promise<boolean> => {
+    try {
+        await api.post('bookmark', { id });
+    } catch (error) {
+        toast({
+            title: "API Error",
+            variant: "destructive",
+            description: `Failed to bookmark result: ${error}`
+        });
+
+        return false;
+    }
+    toast({
+        description: "Result bookmark updated"
+    });
+
+    return true;
+}
+
+export { bookmarkResult };

web/ui/src/lib/api/types.ts

@@ -36,6 +36,7 @@ type galleryResult = {
   screenshot: string;
   failed: boolean;
   technologies: string[];
+  bookmarked: boolean;
 };
 
 // list
@@ -48,6 +49,7 @@ type list = {
   protocol: string;
   content_length: number;
   title: string;
+  bookmarked: boolean;
   failed: boolean;
   failed_reason: string;
 };
@@ -136,6 +138,7 @@ interface detail {
   html: string;
   title: string;
   perception_hash: string;
+  bookmarked: boolean;
   file_name: string;
   is_pdf: boolean;
   failed: boolean;
@@ -159,6 +162,7 @@ interface searchresult {
   matched_fields: string[];
   file_name: string;
   screenshot: string;
+  bookmarked: boolean;
 }
 
 interface technologylist {

web/ui/src/pages/detail/Detail.tsx

@@ -5,6 +5,7 @@ import { Badge } from "@/components/ui/badge";
 import { Button } from "@/components/ui/button";
 import { ScrollArea } from "@/components/ui/scroll-area";
 import { ExternalLink, ChevronLeft, ChevronRight, Code, ClockIcon, Trash2Icon, DownloadIcon, ImagesIcon, ZoomInIcon, CopyIcon } from 'lucide-react';
+import { BookmarkIcon, BookmarkFilledIcon } from "@radix-ui/react-icons";
 import { Dialog, DialogContent, DialogDescription, DialogFooter, DialogHeader, DialogTitle, DialogTrigger, } from "@/components/ui/dialog";
 import { WideSkeleton } from '@/components/loading';
 import { Form, Link, useNavigate, useParams } from 'react-router-dom';
@@ -14,6 +15,7 @@ import { Tabs, TabsContent, TabsList, TabsTrigger } from '@/components/ui/tabs';
 import { copyToClipboard, getIconUrl, getStatusColor } from '@/lib/common';
 import * as api from "@/lib/api/api";
 import * as apitypes from "@/lib/api/types";
+import { bookmarkResult } from "@/lib/api/bookmark";
 import { getData } from './data';
 
 
@@ -93,6 +95,21 @@ const ScreenshotDetailPage = () => {
     }
   };
 
+  const handleBookmarkClick = async () => {
+    const bookmarkUpdated = await bookmarkResult(currentId)
+    if (bookmarkUpdated) {
+      setDetail(prevDetail => {
+         if (!prevDetail) {
+           return prevDetail;
+         }
+         return {
+           ...prevDetail,
+           bookmarked: !prevDetail.bookmarked
+         };
+       });
+    }
+  }
+
   if (loading) return <WideSkeleton />;
   if (!detail) return;
 
@@ -135,6 +152,23 @@ const ScreenshotDetailPage = () => {
               </TooltipContent>
             </Tooltip>
           </TooltipProvider>
+          <TooltipProvider>
+            <Tooltip>
+              <TooltipTrigger asChild>
+                <Button
+                  variant="outline"
+                  size="sm"
+                  onClick={handleBookmarkClick}
+                >
+                  {detail.bookmarked ? <BookmarkFilledIcon className="mr-2 h-4 w-4" />: <BookmarkIcon className="mr-2 h-4 w-4" />}
+                  Bookmark
+                </Button>
+              </TooltipTrigger>
+              <TooltipContent>
+                <p>Bookmark result</p>
+              </TooltipContent>
+            </Tooltip>
+          </TooltipProvider>
           <TooltipProvider>
             <Tooltip>
               <TooltipTrigger asChild>
@@ -719,4 +753,4 @@ const ScreenshotDetailPage = () => {
   );
 };
 
-export default ScreenshotDetailPage;
+export default ScreenshotDetailPage;