rohit/probo
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add mitigation with task and mapping import

Browse Source 
Signed-off-by: gearnode <bryan@frimin.fr>
This commit is contained in:
gearnode
2025-04-01 18:02:43 +02:00
parent 1ee0c3beb7
commit e8dc868edf
13 changed files with 417 additions and 72 deletions
Download Patch File Download Diff File Expand all files Collapse all files

44
pkg/coredata/control.go
View File

@@ -159,6 +159,50 @@ WHERE
return nil
}
func (c *Control) LoadByFrameworkIDAndReferenceID(
ctx context.Context,
conn pg.Conn,
scope Scoper,
frameworkID gid.GID,
referenceID string,
) error {
q := `
SELECT
id,
reference_id,
framework_id,
tenant_id,
name,
description,
created_at,
updated_at
FROM
controls
WHERE
%s
AND framework_id = @framework_id
AND reference_id = @reference_id
LIMIT 1;
`
q = fmt.Sprintf(q, scope.SQLFragment())
args := pgx.StrictNamedArgs{"framework_id": frameworkID, "reference_id": referenceID}
maps.Copy(args, scope.SQLArguments())
rows, err := conn.Query(ctx, q, args)
if err != nil {
return fmt.Errorf("cannot query controls: %w", err)
}
control, err := pgx.CollectExactlyOneRow(rows, pgx.RowToStructByName[Control])
if err != nil {
return fmt.Errorf("cannot collect control: %w", err)
}
*c = control
return nil
}
func (c *Control) LoadByID(
ctx context.Context,
conn pg.Conn,

32
pkg/coredata/control_mitigation.go
View File

@@ -36,6 +36,38 @@ type (
ControlMitigations []*ControlMitigation
)
func (cm ControlMitigation) Upsert(
ctx context.Context,
conn pg.Conn,
scope Scoper,
) error {
q := `
INSERT INTO
controls_mitigations (
control_id,
mitigation_id,
tenant_id,
created_at
)
VALUES (
@control_id,
@mitigation_id,
@tenant_id,
@created_at
)
ON CONFLICT (control_id, mitigation_id) DO NOTHING;
`
args := pgx.StrictNamedArgs{
"control_id": cm.ControlID,
"mitigation_id": cm.MitigationID,
"tenant_id": scope.GetTenantID(),
"created_at": cm.CreatedAt,
}
_, err := conn.Exec(ctx, q, args)
return err
}
func (cm ControlMitigation) Insert(
ctx context.Context,
conn pg.Conn,

42
pkg/coredata/framework.go
View File

@@ -104,6 +104,48 @@ WHERE
return nil
}
func (f *Framework) LoadByReferenceID(
ctx context.Context,
conn pg.Conn,
scope Scoper,
referenceID string,
) error {
q := `
SELECT
id,
organization_id,
reference_id,
name,
description,
created_at,
updated_at
FROM
frameworks
WHERE
%s
AND reference_id = @reference_id
LIMIT 1;
`
q = fmt.Sprintf(q, scope.SQLFragment())
args := pgx.StrictNamedArgs{"reference_id": referenceID}
maps.Copy(args, scope.SQLArguments())
rows, err := conn.Query(ctx, q, args)
if err != nil {
return fmt.Errorf("cannot query frameworks: %w", err)
}
framework, err := pgx.CollectExactlyOneRow(rows, pgx.RowToStructByName[Framework])
if err != nil {
return fmt.Errorf("cannot collect framework: %w", err)
}
*f = framework
return nil
}
func (f *Framework) LoadByID(
ctx context.Context,
conn pg.Conn,

5
pkg/coredata/migrations/20250401T144700Z.sql Normal file
View File

@@ -0,0 +1,5 @@
ALTER TABLE tasks ADD COLUMN reference_id TEXT DEFAULT '';
ALTER TABLE tasks ALTER COLUMN reference_id SET NOT NULL;
ALTER TABLE tasks ALTER COLUMN reference_id DROP DEFAULT;
ALTER TABLE tasks ADD CONSTRAINT tasks_reference_id_unique UNIQUE (id, reference_id);

3
pkg/coredata/migrations/20250401T165000Z.sql Normal file
View File

@@ -0,0 +1,3 @@
ALTER TABLE mitigations RENAME COLUMN content_ref TO reference_id;
ALTER TABLE mitigations ADD CONSTRAINT mitigations_org_ref_unique UNIQUE (organization_id, reference_id);
ALTER TABLE mitigations DROP COLUMN standards;

1
pkg/coredata/migrations/20250401T171500Z.sql Normal file
View File

@@ -0,0 +1 @@
ALTER TABLE mitigations DROP COLUMN version;

11
pkg/coredata/migrations/20250401T173800Z.sql Normal file
View File

@@ -0,0 +1,11 @@
ALTER TABLE controls_mitigations
ADD CONSTRAINT fk_controls_mitigations_control_id
FOREIGN KEY (control_id)
REFERENCES controls(id)
ON DELETE CASCADE;
ALTER TABLE controls_mitigations
ADD CONSTRAINT fk_controls_mitigations_mitigation_id
FOREIGN KEY (mitigation_id)
REFERENCES mitigations(id)
ON DELETE CASCADE;

2
pkg/coredata/migrations/20250401T175300Z.sql Normal file
View File

@@ -0,0 +1,2 @@
ALTER TABLE tasks DROP CONSTRAINT tasks_reference_id_unique;
ALTER TABLE tasks ADD CONSTRAINT tasks_reference_id_unique UNIQUE (mitigation_id, reference_id);

3
pkg/coredata/migrations/20250401T175400Z.sql Normal file
View File

@@ -0,0 +1,3 @@
ALTER TABLE controls
DROP CONSTRAINT IF EXISTS controls_framework_ref_unique,
ADD CONSTRAINT controls_framework_ref_unique UNIQUE (framework_id, reference_id);

127
pkg/coredata/mitigation.go
View File

@@ -37,10 +37,9 @@ type (
Description string `db:"description"`
Importance MitigationImportance `db:"importance"`
State MitigationState `db:"state"`
ContentRef string `db:"content_ref"`
ReferenceID string `db:"reference_id"`
CreatedAt time.Time `db:"created_at"`
UpdatedAt time.Time `db:"updated_at"`
Standards []string `db:"standards"`
}
Mitigations []*Mitigation
@@ -73,10 +72,9 @@ WITH mtgtns AS (
m.description,
m.state,
m.importance,
m.content_ref,
m.reference_id,
m.created_at,
m.updated_at,
m.standards
m.updated_at
FROM
mitigations m
INNER JOIN
@@ -93,10 +91,9 @@ SELECT
description,
state,
importance,
content_ref,
reference_id,
created_at,
updated_at,
standards
updated_at
FROM
mtgtns
WHERE %s
@@ -141,10 +138,9 @@ WITH mtgtns AS (
m.description,
m.state,
m.importance,
m.content_ref,
m.reference_id,
m.created_at,
m.updated_at,
m.standards
m.updated_at
FROM
mitigations m
INNER JOIN
@@ -161,10 +157,9 @@ SELECT
description,
state,
importance,
content_ref,
reference_id,
created_at,
updated_at,
standards
updated_at
FROM
mtgtns
WHERE %s
@@ -208,10 +203,9 @@ SELECT
description,
state,
importance,
content_ref,
reference_id,
created_at,
updated_at,
standards
updated_at
FROM
mitigations
WHERE
@@ -256,10 +250,9 @@ SELECT
description,
state,
importance,
content_ref,
reference_id,
created_at,
updated_at,
standards
updated_at
FROM
mitigations
WHERE
@@ -288,6 +281,87 @@ LIMIT 1;
return nil
}
func (c *Mitigation) Upsert(
ctx context.Context,
conn pg.Conn,
scope Scoper,
) error {
q := `
INSERT INTO
mitigations (
tenant_id,
id,
organization_id,
category,
name,
importance,
state,
description,
reference_id,
created_at,
updated_at
)
VALUES (
@tenant_id,
@mitigation_id,
@organization_id,
@category,
@name,
@importance,
@state,
@description,
@reference_id,
@created_at,
@updated_at
)
ON CONFLICT (organization_id, reference_id) DO UPDATE SET
name = @name,
description = @description,
category = @category,
updated_at = @updated_at
RETURNING
tenant_id,
id,
organization_id,
category,
name,
importance,
state,
description,
reference_id,
created_at,
updated_at
`
args := pgx.StrictNamedArgs{
"tenant_id": scope.GetTenantID(),
"mitigation_id": c.ID,
"organization_id": c.OrganizationID,
"category": c.Category,
"name": c.Name,
"importance": c.Importance,
"state": c.State,
"description": c.Description,
"reference_id": c.ReferenceID,
"created_at": c.CreatedAt,
"updated_at": c.UpdatedAt,
}
rows, err := conn.Query(ctx, q, args)
if err != nil {
return fmt.Errorf("cannot query mitigations: %w", err)
}
mitigation, err := pgx.CollectExactlyOneRow(rows, pgx.RowToStructByName[Mitigation])
if err != nil {
return fmt.Errorf("cannot collect mitigations: %w", err)
}
*c = mitigation
return nil
}
func (c Mitigation) Insert(
ctx context.Context,
conn pg.Conn,
@@ -304,10 +378,9 @@ INSERT INTO
importance,
state,
description,
content_ref,
reference_id,
created_at,
updated_at,
standards
updated_at
)
VALUES (
@tenant_id,
@@ -318,10 +391,9 @@ VALUES (
@importance,
@state,
@description,
@content_ref,
@reference_id,
@created_at,
@updated_at,
@standards
@updated_at
);
`
@@ -332,12 +404,11 @@ VALUES (
"category": c.Category,
"name": c.Name,
"description": c.Description,
"content_ref": c.ContentRef,
"reference_id": c.ReferenceID,
"created_at": c.CreatedAt,
"updated_at": c.UpdatedAt,
"state": c.State,
"importance": c.Importance,
"standards": c.Standards,
}
_, err := conn.Exec(ctx, q, args)
return err

112
pkg/coredata/task.go
View File

@@ -34,11 +34,11 @@ type (
Name string `db:"name"`
Description string `db:"description"`
State TaskState `db:"state"`
ReferenceID string `db:"reference_id"`
TimeEstimate *time.Duration `db:"time_estimate"`
AssignedToID *gid.GID `db:"assigned_to"`
CreatedAt time.Time `db:"created_at"`
UpdatedAt time.Time `db:"updated_at"`
Version int `db:"version"`
}
Tasks []*Task
@@ -66,11 +66,11 @@ SELECT
name,
description,
state,
reference_id,
time_estimate,
assigned_to,
created_at,
updated_at,
version
updated_at
FROM
tasks
WHERE
@@ -112,12 +112,12 @@ INSERT INTO
mitigation_id,
name,
description,
reference_id,
state,
time_estimate,
assigned_to,
created_at,
updated_at,
version
updated_at
)
VALUES (
@tenant_id,
@@ -125,12 +125,12 @@ VALUES (
@mitigation_id,
@name,
@description,
@reference_id,
@state,
@time_estimate,
@assigned_to,
@created_at,
@updated_at,
@version
@updated_at
);
`
@@ -140,17 +140,95 @@ VALUES (
"mitigation_id": c.MitigationID,
"name": c.Name,
"description": c.Description,
"reference_id": c.ReferenceID,
"state": c.State,
"time_estimate": c.TimeEstimate,
"assigned_to": c.AssignedToID,
"created_at": c.CreatedAt,
"updated_at": c.UpdatedAt,
"version": 0,
}
_, err := conn.Exec(ctx, q, args)
return err
}
func (c *Task) Upsert(
ctx context.Context,
conn pg.Conn,
scope Scoper,
) error {
q := `
INSERT INTO
tasks (
tenant_id,
id,
mitigation_id,
name,
description,
reference_id,
state,
time_estimate,
assigned_to,
created_at,
updated_at
)
VALUES (
@tenant_id,
@task_id,
@mitigation_id,
@name,
@description,
@reference_id,
@state,
@time_estimate,
@assigned_to,
@created_at,
@updated_at
)
ON CONFLICT (mitigation_id, reference_id) DO UPDATE SET
name = @name,
description = @description,
updated_at = @updated_at
RETURNING
id,
mitigation_id,
name,
description,
reference_id,
state,
time_estimate,
assigned_to,
created_at,
updated_at
`
args := pgx.StrictNamedArgs{
"tenant_id": scope.GetTenantID(),
"task_id": c.ID,
"mitigation_id": c.MitigationID,
"name": c.Name,
"description": c.Description,
"reference_id": c.ReferenceID,
"state": c.State,
"time_estimate": c.TimeEstimate,
"assigned_to": c.AssignedToID,
"created_at": c.CreatedAt,
"updated_at": c.UpdatedAt,
}
rows, err := conn.Query(ctx, q, args)
if err != nil {
return fmt.Errorf("cannot upsert task: %w", err)
}
task, err := pgx.CollectExactlyOneRow(rows, pgx.RowToStructByName[Task])
if err != nil {
return fmt.Errorf("cannot collect tasks: %w", err)
}
*c = task
return nil
}
func (c *Tasks) LoadByMitigationID(
ctx context.Context,
conn pg.Conn,
@@ -165,11 +243,11 @@ SELECT
name,
description,
state,
reference_id,
time_estimate,
assigned_to,
created_at,
updated_at,
version
updated_at
FROM
tasks
WHERE
@@ -240,8 +318,7 @@ func (c *Task) AssignTo(
q := `
UPDATE tasks SET
assigned_to = @assigned_to,
updated_at = @updated_at,
version = version + 1
updated_at = @updated_at
WHERE %s
AND id = @task_id
RETURNING
@@ -249,12 +326,12 @@ RETURNING
mitigation_id,
name,
description,
reference_id,
state,
time_estimate,
assigned_to,
created_at,
updated_at,
version
updated_at
`
q = fmt.Sprintf(q, scope.SQLFragment())
@@ -289,8 +366,7 @@ func (c *Task) Unassign(
q := `
UPDATE tasks SET
assigned_to = NULL,
updated_at = @updated_at,
version = version + 1
updated_at = @updated_at
WHERE %s
AND id = @task_id
RETURNING
@@ -298,12 +374,12 @@ RETURNING
mitigation_id,
name,
description,
reference_id,
state,
time_estimate,
assigned_to,
created_at,
updated_at,
version
updated_at
`
q = fmt.Sprintf(q, scope.SQLFragment())

105
pkg/probo/mitigation_service.go
View File

@@ -53,6 +53,16 @@ type (
Description string `json:"description"`
Category string `json:"category"`
Importance coredata.MitigationImportance `json:"importance"`
ReferenceID string `json:"reference-id"`
Standards []struct {
Framework string `json:"framework"`
Control string `json:"control"`
} `json:"standards"`
Tasks []struct {
Name string `json:"name"`
Description string `json:"description"`
ReferenceID string `json:"reference-id"`
} `json:"tasks"`
} `json:"mitigations"`
}
)
@@ -124,36 +134,80 @@ func (s MitigationService) Import(
organizationID gid.GID,
req ImportMitigationRequest,
) (*page.Page[*coredata.Mitigation, coredata.MitigationOrderField], error) {
importedMitigations := coredata.Mitigations{}
for _, mitigation := range req.Mitigations {
now := time.Now()
mitigationID, err := gid.NewGID(organizationID.TenantID(), coredata.MitigationEntityType)
if err != nil {
return nil, fmt.Errorf("cannot create global id: %w", err)
}
importedMitigations = append(importedMitigations, &coredata.Mitigation{
ID: mitigationID,
OrganizationID: organizationID,
Name: mitigation.Name,
Description: mitigation.Description,
Category: mitigation.Category,
State: coredata.MitigationStateNotStarted,
Standards: []string{},
Importance: mitigation.Importance,
CreatedAt: now,
UpdatedAt: now,
})
}
err := s.svc.pg.WithTx(
ctx,
func(tx pg.Conn) error {
for _, mitigation := range importedMitigations {
if err := mitigation.Insert(ctx, tx, s.svc.scope); err != nil {
return fmt.Errorf("cannot insert mitigation: %w", err)
for i := range req.Mitigations {
now := time.Now()
mitigationID, err := gid.NewGID(organizationID.TenantID(), coredata.MitigationEntityType)
if err != nil {
return fmt.Errorf("cannot create global id: %w", err)
}
mitigation := &coredata.Mitigation{
ID: mitigationID,
OrganizationID: organizationID,
Name: req.Mitigations[i].Name,
Description: req.Mitigations[i].Description,
Category: req.Mitigations[i].Category,
State: coredata.MitigationStateNotStarted,
ReferenceID: req.Mitigations[i].ReferenceID,
Importance: req.Mitigations[i].Importance,
CreatedAt: now,
UpdatedAt: now,
}
importedMitigations = append(importedMitigations, mitigation)
if err := mitigation.Upsert(ctx, tx, s.svc.scope); err != nil {
return fmt.Errorf("cannot upsert mitigation: %w", err)
}
for j := range req.Mitigations[i].Tasks {
taskID, err := gid.NewGID(organizationID.TenantID(), coredata.TaskEntityType)
if err != nil {
return fmt.Errorf("cannot create global id: %w", err)
}
task := &coredata.Task{
ID: taskID,
MitigationID: mitigation.ID,
Name: req.Mitigations[i].Tasks[j].Name,
Description: req.Mitigations[i].Tasks[j].Description,
ReferenceID: req.Mitigations[i].Tasks[j].ReferenceID,
State: coredata.TaskStateTodo,
CreatedAt: now,
UpdatedAt: now,
}
if err := task.Upsert(ctx, tx, s.svc.scope); err != nil {
return fmt.Errorf("cannot upsert task: %w", err)
}
}
for _, standard := range req.Mitigations[i].Standards {
framework := &coredata.Framework{}
if err := framework.LoadByReferenceID(ctx, tx, s.svc.scope, standard.Framework); err != nil {
continue
}
control := &coredata.Control{}
if err := control.LoadByFrameworkIDAndReferenceID(ctx, tx, s.svc.scope, framework.ID, standard.Control); err != nil {
continue
}
controlMitigation := &coredata.ControlMitigation{
ControlID: control.ID,
MitigationID: mitigation.ID,
CreatedAt: now,
}
if err := controlMitigation.Upsert(ctx, tx, s.svc.scope); err != nil {
return fmt.Errorf("cannot insert control mitigation: %w", err)
}
}
}
@@ -271,7 +325,6 @@ func (s MitigationService) Create(
Description: req.Description,
Category: req.Category,
State: coredata.MitigationStateNotStarted,
Standards: []string{},
Importance: req.Importance,
CreatedAt: now,
UpdatedAt: now,

2
pkg/probo/task_service.go
View File

@@ -22,6 +22,7 @@ import (
"github.com/getprobo/probo/pkg/coredata"
"github.com/getprobo/probo/pkg/gid"
"github.com/getprobo/probo/pkg/page"
"github.com/getprobo/probo/pkg/slug"
"go.gearno.de/kit/pg"
)
@@ -65,6 +66,7 @@ func (s TaskService) Create(
TimeEstimate: req.TimeEstimate,
AssignedToID: req.AssignedToID,
State: coredata.TaskStateTodo,
ReferenceID: slug.Make(req.Name),
CreatedAt: now,
UpdatedAt: now,
}