feat: initial access check (#104)

* feat: enhance command state classification for initial access detection

* fix: improve prompt clarity for initial access detection

* docs: update graph visualization

* feat: enhance command output classification with detailed priority order for initial access detection

* fix: clarify command state classification rules and enforce error handling for unknown states

* docs: update graph visualization

* fix: improve prompt clarity in command state classifier by formatting multiline instructions

* docs: update graph visualization

---------

Co-authored-by: GitHub Action <action@github.com>

Author: laysakura

wish-command-generation-api/docs/graph.svg

@@ -12,19 +12,77 @@
 
 # Define the prompt template
 COMMAND_STATE_CLASSIFIER_PROMPT = """
-As a system operations analyst, your role is to analyze command outputs and report the results.
-Follow these specific steps:
+As a security operations analyst, your task is to determine if a command output indicates successful
+initial access to a target system.
+
+# PRIORITY ORDER FOR CLASSIFICATION:
+
+1. FIRST: Check for Initial Access Success
+   - If ANY initial access indicators are found, output "SUCCESS_INITIAL_ACCESS" regardless of exit_code
+   - This takes absolute precedence over all other classifications
+
+2. SECOND: Check for Normal Success
+   - Only if NO initial access indicators were found AND exit_code is "0", output "SUCCESS"
+   - Even if exit_code is 0, do not output "SUCCESS" if there are initial access indicators
+     (use "SUCCESS_INITIAL_ACCESS" instead)
+
+3. THIRD: Check for Specific Error Types
+   - Only if neither of the above conditions are met
+
+# MOST IMPORTANT RULE: Check for Initial Access Success REGARDLESS of exit_code
+
+Even if the command timed out (exit_code 124) or had any other error, FIRST check if there are ANY signs
+of successful initial access in the output.
+
+# FIRST PRIORITY: Check for Initial Access Success
+
+Before any other analysis, check if the command output indicates successful initial access to the target system.
+Output "SUCCESS_INITIAL_ACCESS" if ANY of the following indicators are present:
+
+## 1. Meterpreter or shell session establishment (HIGHEST PRIORITY)
+- ANY line containing "Meterpreter session" AND "opened" (This is the strongest indicator)
+- ANY line containing "Session" AND "created"
+- ANY line containing "Command shell session" AND "opened"
+- ANY line containing "WIN" surrounded by special characters (like =-=-=WIN=-=-=)
+- "Successful connection established"
+- "Successfully opened reverse shell"
+
+## 2. Successful remote command execution proving access
+- Evidence of successful `whoami` execution on target
+- Evidence of successful `id` execution on target
+- Evidence of successful `hostname` execution on target
+- Evidence of successful `ipconfig`/`ifconfig` execution on target
+- Evidence of successful `ls`/`dir` execution on target
+
+## 3. Explicit access confirmation
+- "Access gained to target system"
+- "Successfully compromised target"
+- "Initial foothold established"
+- "Initial access achieved"
+- "Successfully exploited vulnerability"
+
+# IMPORTANT: The exit_code does NOT override initial access indicators
+If ANY of the above indicators are found, output "SUCCESS_INITIAL_ACCESS" regardless of the exit_code.
+This includes cases where the command timed out (exit_code 124) but still managed to establish a session.
+
+# SECOND PRIORITY: Standard Command State Analysis
+
+Only if NONE of the above initial access indicators are found, proceed with standard command state analysis:
 
 1. If the `exit_code` is "0", output "SUCCESS" and end.
+   - This includes successful port scans that found open ports
+   - This includes successful reconnaissance commands
+   - This includes ANY command that completed normally with exit_code 0
+
 2. Otherwise, check the command output from `stdout` and `stderr`. Then choose the most appropriate
-   error code from the following:
+   code from the following:
    - TIMEOUT: When command execution times out or when the output indicates a timeout occurred.
      Look for indicators such as:
      * Messages containing phrases like "timeout", "timed out", "Read timeout expired"
      * Messages containing "connection timeout" or similar connection-related timeout phrases
      * Log entries indicating that a command was terminated due to exceeding time limits
      * Warning messages about waiting too long for command completion
-     * `exit_code` is "-1" (reserved for timeout)
+     * `exit_code` is "124" or "-1" (reserved for timeout)
    - COMMAND_NOT_FOUND: When the command is not found on the local machine
    - FILE_NOT_FOUND: When a local file referenced in the command is not found (excluding remote files like smb)
    - REMOTE_OPERATION_FAILED: When an operation on a remote machine (e.g., file reference, command execution) fails
@@ -33,8 +91,11 @@
 
 3. Output the selected error code and end.
 
-Note that OTHERS should only be used when the error does not fit any of the above categories
-and there are no success indicators in the output.
+# NOTE
+
+- SUCCESS_INITIAL_ACCESS takes precedence over all other states, including SUCCESS and TIMEOUT.
+  OTHERS is the last resort.
+- Output only the classification string, without any extra characters.
 
 # command
 {command}
@@ -99,7 +160,9 @@ def classify_command_state(state: GraphState, settings_obj: Settings) -> GraphSt
         ).strip()
 
         # Convert the classification string to CommandState
-        if classification_str == "SUCCESS":
+        if classification_str == "SUCCESS_INITIAL_ACCESS":
+            command_state = CommandState.SUCCESS_INITIAL_ACCESS
+        elif classification_str == "SUCCESS":
             command_state = CommandState.SUCCESS
         elif classification_str == "COMMAND_NOT_FOUND":
             command_state = CommandState.COMMAND_NOT_FOUND
@@ -111,8 +174,10 @@ def classify_command_state(state: GraphState, settings_obj: Settings) -> GraphSt
             command_state = CommandState.TIMEOUT
         elif classification_str == "NETWORK_ERROR":
             command_state = CommandState.NETWORK_ERROR
-        else:
+        elif classification_str == "OTHERS":
             command_state = CommandState.OTHERS
+        else:
+            raise ValueError(f"Unknown command state classification: {classification_str}")
 
         # Set the command state in the new state
         new_state.command_state = command_state

wish-log-analysis-api/docs/graph.svg

@@ -7,6 +7,11 @@ class CommandState(str, Enum):
     DOING = "DOING"
     """The command is being executed."""
 
+    SUCCESS_INITIAL_ACCESS = "SUCCESS_INITIAL_ACCESS"
+    """The command succeeded with initial access to the target system.
+
+    Examples include establishing a Meterpreter session or getting a shell."""
+
     SUCCESS = "SUCCESS"
     """The command succeeded, with meaningful insights."""