From 14743745d4d20e94c9924831b6b2153f7bab10de Mon Sep 17 00:00:00 2001
From: nlf <nlf@nlf.sh>
Date: Thu, 19 Oct 2023 16:25:29 -0700
Subject: [PATCH] feat: get probability parsing added, for #6

---
 lib/parser/helpers.ts   | 35 +++++++++++++++++++++++++++++++++--
 lib/parser/tree.peggy   | 10 ++++++----
 test/fixtures/sample.td |  2 +-
 3 files changed, 40 insertions(+), 7 deletions(-)

diff --git a/lib/parser/helpers.ts b/lib/parser/helpers.ts
index 8c18f81..358545b 100644
--- a/lib/parser/helpers.ts
+++ b/lib/parser/helpers.ts
@@ -22,6 +22,30 @@ export function parseMitigation(mitigation?: MitigationInput): MitigationOutput
   return result;
 }
 
+export type ProbabilityInput = "(high)" | "(H)" | ("medium") | "(M)" | "(low)" | "(L)";
+export interface ProbabilityOutput {
+  probability: "high" | "medium" | "low";
+}
+export function parseProbability(probability?: ProbabilityInput): ProbabilityOutput | undefined {
+  if (!probability) {
+    return undefined;
+  }
+
+  if (["(high)", "(H)"].includes(probability)) {
+    return { probability: "high" };
+  }
+
+  if (["(medium)", "(M)"].includes(probability)) {
+    return { probability: "medium" };
+  }
+
+  if (["(low)", "(L)"].includes(probability)) {
+    return { probability: "low" };
+  }
+
+  return undefined;
+}
+
 export interface RawChild {
   and?: Child;
   or?: Child;
@@ -63,12 +87,14 @@ export interface ChildInput {
   boolean: BooleanInput;
   children?: RawChild[];
   mitigation?: MitigationInput;
+  probability?: ProbabilityInput;
 }
 
 export interface Assumption {
   assumption: string;
   comments?: string[];
   mitigation?: boolean;
+  probability?: "high" | "medium" | "low";
   complete?: boolean;
 }
 
@@ -78,16 +104,18 @@ export interface AssumptionChild {
   or?: Assumption;
 }
 
-export function createAssumption({ line, boolean, mitigation, children }: ChildInput): AssumptionChild {
+export function createAssumption({ line, boolean, mitigation, probability, children }: ChildInput): AssumptionChild {
   const parsedLine = parseLine(line);
   const parsedBoolean = parseBoolean(boolean);
   const parsedMitigation = parseMitigation(mitigation);
+  const parsedProbability = parseProbability(probability);
   const parsedChildren = parseChildren(children);
 
   const result: AssumptionChild = {
     [parsedBoolean]: {
       assumption: parsedLine,
       ...parsedMitigation,
+      ...parsedProbability,
       ...parsedChildren,
     },
   };
@@ -99,6 +127,7 @@ export interface Condition {
   condition: string;
   comments?: string[];
   mitigation?: boolean;
+  probability?: "high" | "medium" | "low";
   complete?: boolean;
 }
 
@@ -108,16 +137,18 @@ export interface ConditionChild {
   or?: Condition;
 }
 
-export function createCondition({ line, boolean, mitigation, children }: ChildInput): ConditionChild {
+export function createCondition({ line, boolean, mitigation, probability, children }: ChildInput): ConditionChild {
   const parsedLine = parseLine(line);
   const parsedBoolean = parseBoolean(boolean);
   const parsedMitigation = parseMitigation(mitigation);
+  const parsedProbability = parseProbability(probability);
   const parsedChildren = parseChildren(children);
 
   return {
     [parsedBoolean]: {
       condition: parsedLine,
       ...parsedMitigation,
+      ...parsedProbability,
       ...parsedChildren,
     },
   };
diff --git a/lib/parser/tree.peggy b/lib/parser/tree.peggy
index bdc7a5e..5fbafbc 100644
--- a/lib/parser/tree.peggy
+++ b/lib/parser/tree.peggy
@@ -28,14 +28,14 @@ Objective =
   { return createObjective({ line, children }) }
 
 Assumption =
-  SAME_INDENT boolean:Boolean "?" WS mitigation:Mitigation? line:LineToEnd EOL?
+  SAME_INDENT boolean:Boolean "?" WS mitigation:Mitigation? WS probability:Probability? line:LineToEnd EOL?
   children:IndentedChildren? EOL*
-  { return createAssumption({ boolean, mitigation, line, children }) }
+  { return createAssumption({ boolean, mitigation, probability, line, children }) }
 
 Condition =
-  SAME_INDENT boolean:Boolean WS mitigation:Mitigation? line:LineToEnd EOL?
+  SAME_INDENT boolean:Boolean WS mitigation:Mitigation? WS probability:Probability? line:LineToEnd EOL?
   children:IndentedChildren? EOL*
-  { return createCondition({ boolean, mitigation, line, children }) }
+  { return createCondition({ boolean, mitigation, probability, line, children }) }
 
 Comment =
   SAME_INDENT ">" WS comment:LineToEnd EOL?
@@ -55,6 +55,8 @@ Boolean = "-" / "+"
 
 Mitigation = "[ ]" / "[x]" / "[X]"
 
+Probability = "(high)" / "(H)" / "(medium)" / "(M)" / "(low)" / "(L)"
+
 LineToEnd = (!EOL char:. { return char })+
 
 WS = [ \t]*
diff --git a/test/fixtures/sample.td b/test/fixtures/sample.td
index cb24f23..3ad31ad 100644
--- a/test/fixtures/sample.td
+++ b/test/fixtures/sample.td
@@ -1,6 +1,6 @@
 __Attacker's goal__
   - method which in order to be viable
-    + requires this condition to be true
+    + (high) requires this condition to be true
     + and this condition which depends on either
       - x to be true
       - or y to be true