Make risk magnitude a struct (#1352)

mondoohq · Jul 12, 2024 · 4ca0be1 · 4ca0be1
1 parent 884152d
commit 4ca0be1
Show file tree

Hide file tree

Showing 9 changed files with 1,906 additions and 1,688 deletions.
diff --git a/internal/bundle/bundle.yac.go b/internal/bundle/bundle.yac.go
diff --git a/internal/bundle/bundle_ext.go b/internal/bundle/bundle_ext.go
@@ -150,3 +150,19 @@ func (x *Remediation) MarshalYAML() (interface{}, error) {
 
  return x.Items, nil
 }
+
+func (x *RiskMagnitude) UnmarshalYAML(node *yaml.Node) error {
+ x.addFileContext(node)
+
+ var res float32
+ if err := node.Decode(&res); err == nil {
+ x.Value = res
+ return nil
+ }
+
+ type tmp RiskMagnitude
+ if err := node.Decode((*tmp)(x)); err != nil {
+ return errors.Wrap(err, "can't unmarshal risk magnitude")
+ }
+ return nil
+}
diff --git a/internal/bundle/yacit/main.go b/internal/bundle/yacit/main.go
@@ -17,7 +17,7 @@ import (
 func main() {
  res := yacit.New(yacit.YacItConfig{
  SkipUnmarshal: []string{
- "Impact", "ImpactValue", "Filters", "Remediation",
+ "Impact", "ImpactValue", "Filters", "Remediation", "RiskMagnitude",
  },
  Package: "bundle",
  // field names with sort weight

diff --git a/internal/datalakes/inmemory/policyresolver.go b/internal/datalakes/inmemory/policyresolver.go
@@ -539,8 +539,8 @@ func (db *Db) GetScoredRisks(ctx context.Context, assetMrn string) (*policy.Scor
  return nil, errors.New("cannot find risk metadata for " + risk.Mrn)
  }
  riskInfo := raw.(*policy.RiskFactor)
- srisk.IsAbsolute = riskInfo.IsAbsolute
- srisk.Risk = riskInfo.Magnitude
+ srisk.IsToxic = riskInfo.GetMagnitude().GetIsToxic()
+ srisk.Risk = riskInfo.GetMagnitude().GetValue()
  if !risk.IsDetected {
  srisk.Risk *= -1
  }