fix sort for indirect attributes

packages/kestrel_core/src/kestrel/codegen/commands.py

@@ -20,6 +20,7 @@
 
 import functools
 import logging
+import re
 import itertools
 from collections import OrderedDict
 
@@ -439,13 +440,9 @@ def group(stmt, session):
 @_default_output
 @_skip_command_if_empty_input
 def sort(stmt, session):
-    session.store.assign(
-        stmt["output"],
-        session.symtable[stmt["input"]].entity_table,
-        op="sort",
-        by=stmt["attribute"],
-        ascending=stmt["ascending"],
-    )
+    entity_table = session.symtable[stmt["input"]].entity_table
+    qry = _build_query(session.store, entity_table, Query(entity_table), stmt, [])
+    session.store.assign_query(stmt["output"], qry)
 
 
 @_debug_logger
@@ -532,6 +529,13 @@ def _build_query(store, entity_table, qry, stmt, paths=None):
     if sort_by:
         direction = "ASC" if stmt["ascending"] else "DESC"
         qry.append(Order([(sort_by, direction)]))
+    else:
+        # Check if we need to preserve original sort order
+        # This is kind of a hack, copied from firepit.
+        viewdef = store._get_view_def(entity_table)
+        match = re.search(r"ORDER BY \"([a-z0-9:'\._\-]*)\" (ASC|DESC)$", viewdef)
+        if match:
+            qry.append(Order([(match.group(1), match.group(2))]))
     limit = stmt.get("limit")
     if limit:
         qry.append(Limit(limit))

packages/kestrel_core/tests/test_command_assign.py

@@ -99,3 +99,10 @@ def test_assign_from_empty_var(proc_bundle_file):
         q = s.symtable["q"]
         assert len(q) == 0
         assert q.records_count == 0
+
+
+def test_assign_sort_by_indirect_attr(proc_bundle_file):
+    with Session() as s:
+        s.execute(f"c = GET network-traffic FROM file://{proc_bundle_file} WHERE src_port > 0")
+        s.execute("d = c SORT BY dst_ref.value")
+        _ = s.get_variable("d")

packages/kestrel_core/tests/test_command_disp.py

@@ -96,3 +96,15 @@ def test_disp_empty_variable():
         s.execute(stmt)
         out = s.execute("DISP v2")
         assert out == []
+
+
+def test_disp_sorted_indirect_attr(proc_bundle_file):
+    with Session() as s:
+        s.execute("p = GET process"
+                  f"   FROM file://{proc_bundle_file}"
+                  "    WHERE name in ('conhost.exe', 'services.exe', 'taskhost.exe')")
+        s.execute("sp = SORT p BY binary_ref.name DESC")
+        out = s.execute("DISP sp")
+        data = out[0].to_dict()["data"]
+        names = [d["binary_ref.name"] for d in data]
+        assert names == sorted(names, reverse=True)