added quality of life changes

fslaktern · Nov 17, 2023 · e89b286 · e89b286
1 parent 49c96e3
commit e89b286
Show file tree

Hide file tree

Showing 8 changed files with 63 additions and 26 deletions.
diff --git a/docker-compose-test.yml b/docker-compose-test.yml
@@ -9,6 +9,7 @@ services:
       - "27017:27017"
 
   frontend:
+    restart: always
     build:
       context: frontend
       dockerfile: Dockerfile-frontend
@@ -24,6 +25,7 @@ services:
       API_SERVER_ENDPOINT: http://api:5000/
 
   api:
+    restart: always
     build:
       context: .
       dockerfile: Dockerfile-python
@@ -46,6 +48,7 @@ services:
 
   # Only for testing
   flagidendpoint:
+    restart: always
     build:
       context: services/go-importer/test_data
     image: flagid-endpoint:latest
@@ -57,6 +60,7 @@ services:
       - "8000:8000"
 
   flagids:
+    restart: always
     build:
       context: services/flagids
     image: tulip-flagids:latest
@@ -72,6 +76,7 @@ services:
       FLAGID_ENDPOINT: ${FLAGID_ENDPOINT}
 
   assembler:
+    restart: always
     build:
       context: services/go-importer
       dockerfile: Dockerfile-assembler
@@ -92,6 +97,7 @@ services:
       DELAY: 30
 
   enricher:
+    restart: always
     build:
       context: services/go-importer
       dockerfile: Dockerfile-enricher

diff --git a/docker-compose.yml b/docker-compose.yml
@@ -9,6 +9,7 @@ services:
       - "27017:27017"
 
   frontend:
+    restart: always
     build:
       context: frontend
       dockerfile: Dockerfile-frontend
@@ -24,6 +25,7 @@ services:
       API_SERVER_ENDPOINT: http://api:5000/
 
   api:
+    restart: always
     build:
       context: .
       dockerfile: Dockerfile-python
@@ -45,6 +47,7 @@ services:
       VM_IP: ${VM_IP}
 
   flagids:
+    restart: always
     build:
       context: services/flagids
     image: tulip-flagids:latest
@@ -60,6 +63,7 @@ services:
       FLAGID_ENDPOINT: ${FLAGID_ENDPOINT}
 
   assembler:
+    restart: always
     build:
       context: services/go-importer
       dockerfile: Dockerfile-assembler
@@ -78,6 +82,7 @@ services:
       FLAG_LIFETIME: ${FLAG_LIFETIME}
 
   enricher:
+    restart: always
     build:
       context: services/go-importer
       dockerfile: Dockerfile-enricher

diff --git a/frontend/src/pages/FlowView.tsx b/frontend/src/pages/FlowView.tsx
@@ -17,7 +17,10 @@ import {
   useLazyToPwnToolsQuery,
   useToSinglePythonRequestQuery,
 } from "../api";
-import { API_BASE_PATH } from "../const";
+import {
+  API_BASE_PATH,
+  TEXT_FILTER_KEY
+} from "../const";
 
 const SECONDARY_NAVBAR_HEIGHT = 50;
 
@@ -190,6 +193,8 @@ function formatIP(ip: string) {
 }
 
 function FlowOverview({ flow }: { flow: FullFlow }) {
+  const FILTER_KEY = TEXT_FILTER_KEY;
+  let [searchParams, setSearchParams] = useSearchParams();
   return (
     <div>
       {flow.signatures?.length > 0 ? (
@@ -242,26 +247,32 @@ function FlowOverview({ flow }: { flow: FullFlow }) {
             [{flow.flags.map((query, i) => (
             <span>
               {i > 0 ? ', ' : ''}
-              <a
-                  key={query}
-                  href={`/?text=${encodeURIComponent(query)}`}
+              <button className="font-bold"
+                  onClick={() => {
+                    searchParams.set(FILTER_KEY, query);
+                    setSearchParams(searchParams);
+                  }
+                }
               >
               {query}
-              </a>
+              </button>
             </span>
             ))}]
           </div>
           <div>Flagids: </div>
           <div className="font-bold">
             [{flow.flagids.map((query, i) => (
-                <span>
+              <span>
                 {i > 0 ? ', ' : ''}
-                  <a
-                      key={query}
-                      href={`/?text=${encodeURIComponent(query)}`}
-                  >
-                {query}
-              </a>
+                <button className="font-bold"
+                  onClick={() => {
+                      searchParams.set(FILTER_KEY, query);
+                      setSearchParams(searchParams);
+                    }
+                  }
+                >
+                  {query}
+                </button>
               </span>
             ))}]
           </div>

diff --git a/services/flagids/Dockerfile b/services/flagids/Dockerfile
@@ -2,9 +2,12 @@ FROM python:3.8
 
 WORKDIR /app
 
+COPY ./ /app/
+
 COPY ./requirements.txt /app/
 
 RUN pip install -r ./requirements.txt
 
 COPY ./flagids.py /app/
+
 CMD python3 ./flagids.py
diff --git a/services/flagids/flagids.py b/services/flagids/flagids.py
@@ -11,6 +11,8 @@
 start_date = os.getenv("TICK_START", "2018-06-27T13:00+02:00")
 mongo_host = os.getenv("TULIP_MONGO", "localhost:27017").split(':')
 team_id = os.getenv("TEAM_ID", "10.10.3.1")
+team_id_is_digit = team_id.isdigit()
+team_id_int = int(team_id) if team_id_is_digit else None
 flagid_endpoint = os.getenv("FLAGID_ENDPOINT", "http://localhost:8000/flagids.json")
 
 print('STARTING FLAGIDS')
@@ -23,22 +25,28 @@ def get_leaf_nodes(data):
     if isinstance(data, dict):
         if team_id in data.keys():
             yield from get_leaf_nodes(data[team_id])
+        elif team_id_is_digit and team_id_int in data.keys():
+            yield from get_leaf_nodes(data[team_id_int])
         else:
             for key, value in data.items():
                 yield from get_leaf_nodes(value)
     elif isinstance(data, list):
-        for item in data:
-            yield from get_leaf_nodes(item)
+        if team_id in data or (team_id_is_digit and team_id_int in data):
+            yield
+        else:
+            for item in data:
+                print(item, end=' ')
+                yield from get_leaf_nodes(item)
     else:
         # prevent id from being used as Flagids
-        if data != team_id:
-            yield data
+        yield data
 
 def update_flagids():
     print('Updating flagids: ', time.time())
     response = requests.get(flagid_endpoint)
     crnt_time = int(time.time())
-    nodes = [{"_id": node, "time": crnt_time} for node in  get_leaf_nodes(response.json())]
+    nodes = [{"_id": node, "time": crnt_time} for node in get_leaf_nodes(response.json()) if node is not None]
+    print(nodes)
     db['flagids'].insert_many(nodes)
 
 def main():

diff --git a/services/go-importer/cmd/assembler/main.go b/services/go-importer/cmd/assembler/main.go
@@ -61,12 +61,12 @@ func reassemblyCallback(entry db.FlowEntry) {
 	}
 
 	//Apply flagid
-	flagids, err := g_db.GetFlagids()
+	flagids, err := g_db.GetFlagids(flaglifetime)
 	if err != nil {
 		log.Fatal(err)
 	}
 
-	ApplyFlagids(&entry, flagids, flaglifetime)
+	ApplyFlagids(&entry, flagids)
 
 	// Finally, insert the new entry
 	g_db.InsertFlow(entry)

diff --git a/services/go-importer/cmd/assembler/tags.go b/services/go-importer/cmd/assembler/tags.go
@@ -71,16 +71,13 @@ func ApplyFlagTags(flow *db.FlowEntry, reg *string) {
 
 // Apply flagids to the entire flow.
 // This assumes the `Data` part of the flowItem is already pre-processed, s.t.
-func ApplyFlagids(flow *db.FlowEntry, flagids []db.Flagid, flaglifetime int) {
+func ApplyFlagids(flow *db.FlowEntry, flagids []db.Flagid) {
 
 	for idx := 0; idx < len(flow.Flow); idx++ {
 		flowItem := &flow.Flow[idx]
 		data := flowItem.Data
 		for _, flagid := range flagids {
 			flagidstr := flagid.ID
-			if flaglifetime > 0 && (flowItem.Time+flaglifetime) > flagid.Time {
-				continue
-			}
 			if strings.Contains(data, flagidstr) {
 				var tag string
 				if flowItem.From == "c" {

diff --git a/services/go-importer/internal/pkg/db/db.go b/services/go-importer/internal/pkg/db/db.go
@@ -317,16 +317,23 @@ type Flagid struct {
 	Time int    `bson:"time"`
 }
 
-func (db Database) GetFlagids() ([]Flagid, error) {
+func (db Database) GetFlagids(flaglifetime int) ([]Flagid, error) {
 	// Create a context with a timeout
 	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
 	defer cancel()
 
 	// Access the "pcap" database and "flagids" collection
 	collection := db.client.Database("pcap").Collection("flagids")
 
-	// Find all documents in the collection
-	cur, err := collection.Find(ctx, bson.M{})
+	// Find all documents in the
+	var filter bson.M
+	if flaglifetime < 0 {
+		filter = bson.M{}
+	} else {
+		filter = bson.M{"time": bson.M{"$gt": int(time.Now().Unix()) - flaglifetime}}
+	}
+
+	cur, err := collection.Find(ctx, filter)
 	if err != nil {
 		return nil, err
 	}