Skip to content

Commit

Permalink
Updated plugin generator to use SCI types
Browse files Browse the repository at this point in the history 
Signed-off-by: Eddie Knight <[email protected]>
  • Loading branch information
@eddie-knight
eddie-knight committed Feb 23, 2025
1 parent 03d6eb0 commit 87f5cc7
Show file tree
Hide file tree
Showing 3 changed files with 37 additions and 141 deletions.
173 changes: 34 additions & 139 deletions cmd/generate-plugin.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,89 +3,21 @@ package cmd
import (
"fmt"
"html/template"
"net/http"
"os"
"path/filepath"
"strings"

"github.com/go-git/go-git/v5"
"github.com/spf13/cobra"
"github.com/spf13/viper"
"gopkg.in/yaml.v2"
)

type ControlCatalog struct {
CategoryIDFriendly string
ServiceName string
TestSuites map[string][]string

Metadata Metadata `yaml:"metadata"`

Controls []Control `yaml:"controls"`
Features []Feature `yaml:"features"`
Threats []Threat `yaml:"threats"`

LatestReleaseDetails ReleaseDetails `yaml:"latest_release_details"`
}

// Metadata is a struct that represents the metadata.yaml file
type Metadata struct {
Title string `yaml:"title"`
ID string `yaml:"id"`
Description string `yaml:"description"`
ReleaseDetails []ReleaseDetails `yaml:"release_details"`
}

type ReleaseDetails struct {
Version string `yaml:"version"`
AssuranceLevel string `yaml:"assurance_level"`
ThreatModelURL string `yaml:"threat_model_url"`
ThreatModelAuthor string `yaml:"threat_model_author"`
RedTeam string `yaml:"red_team"`
RedTeamExerciseURL string `yaml:"red_team_exercise_url"`
ReleaseManager ReleaseManager `yaml:"release_manager"`
ChangeLog []string `yaml:"change_log"`
}

type ReleaseManager struct {
Name string `yaml:"name"`
GithubId string `yaml:"github_id"`
Company string `yaml:"company"`
Summary string `yaml:"summary"`
}

type Feature struct {
ID string `yaml:"id"`
Title string `yaml:"title"`
Description string `yaml:"description"`
}

type Threat struct {
ID string `yaml:"id"`
Title string `yaml:"title"`
Description string `yaml:"description"`
Features []string `yaml:"features"`
MITRE []string `yaml:"mitre_attack"`
}

type Control struct {
IDFriendly string
ID string `yaml:"id"`
Title string `yaml:"title"`
Objective string `yaml:"objective"`
ControlFamily string `yaml:"control_family"`
Threats []string `yaml:"threats"`
NISTCSF string `yaml:"nist_csf"`
MITREATTACK string `yaml:"mitre_attack"`
ControlMappings map[string]interface{} `yaml:"control_mappings"`
TestRequirements []TestRequirement `yaml:"test_requirements"`
}
"github.com/revanite-io/sci/pkg/layer2"
)

type TestRequirement struct {
IDFriendly string
ID string `yaml:"id"`
Text string `yaml:"text"`
TLPLevels []string `yaml:"tlp_levels"`
type CatalogData struct {
layer2.Catalog
ServiceName string
TestSuites map[string][]string
}

var TemplatesDir string
Expand All @@ -102,8 +34,6 @@ var genPluginCmd = &cobra.Command{
}

func init() {
rootCmd.AddCommand(genPluginCmd)

genPluginCmd.PersistentFlags().StringP("source-path", "p", "", "The source file to generate the plugin from.")
genPluginCmd.PersistentFlags().StringP("local-templates", "", "", "Path to a directory to use instead of downloading the latest templates.")
genPluginCmd.PersistentFlags().StringP("service-name", "n", "", "The name of the service (e.g. 'ECS, AKS, GCS').")
Expand All @@ -113,6 +43,8 @@ func init() {
viper.BindPFlag("local-templates", genPluginCmd.PersistentFlags().Lookup("local-templates"))
viper.BindPFlag("service-name", genPluginCmd.PersistentFlags().Lookup("service-name"))
viper.BindPFlag("output-dir", genPluginCmd.PersistentFlags().Lookup("output-dir"))

rootCmd.AddCommand(genPluginCmd)
}

func generatePlugin() {
Expand Down Expand Up @@ -189,8 +121,18 @@ func setupTemplatesDir() error {
return err
}

func generateFileFromTemplate(data ControlCatalog, templatePath, OutputDir string) error {
tmpl, err := template.ParseFiles(templatePath)
func generateFileFromTemplate(data CatalogData, templatePath, OutputDir string) error {
templateContent, err := os.ReadFile(templatePath)
if err != nil {
return fmt.Errorf("error reading template file %s: %w", templatePath, err)
}

tmpl, err := template.New("plugin").Funcs(template.FuncMap{
"prettify": func(s string) string {
return strings.TrimSpace(strings.ReplaceAll(
strings.ReplaceAll(s, "\n", " "), ".", "_"))
},
}).Parse(string(templateContent))
if err != nil {
return fmt.Errorf("error parsing template file %s: %w", templatePath, err)
}
Expand All @@ -201,14 +143,17 @@ func generateFileFromTemplate(data ControlCatalog, templatePath, OutputDir strin
}

outputPath := filepath.Join(OutputDir, strings.TrimSuffix(relativePath, ".txt"))

err = os.MkdirAll(filepath.Dir(outputPath), os.ModePerm)
if err != nil {
return fmt.Errorf("error creating directories for %s: %w", outputPath, err)
}

outputFile, err := os.Create(outputPath)
if err != nil {
return fmt.Errorf("error creating output file %s: %w", outputPath, err)
}

defer outputFile.Close()

err = tmpl.Execute(outputFile, data)
Expand All @@ -219,76 +164,26 @@ func generateFileFromTemplate(data ControlCatalog, templatePath, OutputDir strin
return nil
}

func readData() (data ControlCatalog, err error) {
if strings.HasPrefix(SourcePath, "http") {
data, err = readYAMLURL()
} else {
data, err = readYAMLFile()
}
func readData() (data CatalogData, err error) {
err = data.LoadControlFamiliesFile(SourcePath)
if err != nil {
return
}

data.TestSuites = make(map[string][]string)
data.CategoryIDFriendly = strings.ReplaceAll(data.Metadata.ID, ".", "_")

for i := range data.Controls {
fmt.Println(data.Controls[i].ID)
data.Controls[i].IDFriendly = strings.ReplaceAll(data.Controls[i].ID, ".", "_")
// loop over objectives in test_requirements and replace newlines with empty string
for j, testReq := range data.Controls[i].TestRequirements {
// Some test requirements have newlines in them, which breaks the template
data.Controls[i].TestRequirements[j].Text = strings.TrimSpace(strings.ReplaceAll(testReq.Text, "\n", " "))
// Replace periods with underscores for the friendly ID
data.Controls[i].TestRequirements[j].IDFriendly = strings.ReplaceAll(testReq.ID, ".", "_")

// Add the test ID to the TestSuites map for each TLP level
for _, tlpLevel := range testReq.TLPLevels {
if data.TestSuites[tlpLevel] == nil {
data.TestSuites[tlpLevel] = []string{}
for i, family := range data.ControlFamilies {
for j := range family.Controls {
for _, testReq := range data.ControlFamilies[i].Controls[j].Requirements {
// Add the test ID to the TestSuites map for each TLP level
for _, tlpLevel := range testReq.Applicability {
if data.TestSuites[tlpLevel] == nil {
data.TestSuites[tlpLevel] = []string{}
}
data.TestSuites[tlpLevel] = append(data.TestSuites[tlpLevel], testReq.ID)
}
data.TestSuites[tlpLevel] = append(data.TestSuites[tlpLevel], strings.ReplaceAll(testReq.ID, ".", "_"))
}
}
}
return
}

func readYAMLURL() (data ControlCatalog, err error) {
resp, err := http.Get(SourcePath)
if err != nil {
logger.Error("Failed to fetch URL: %v", err)
return
}
defer resp.Body.Close()

if resp.StatusCode != http.StatusOK {
logger.Error("Failed to fetch URL: %v", resp.Status)
return
}

decoder := yaml.NewDecoder(resp.Body)
err = decoder.Decode(&data)
if err != nil {
logger.Error("Failed to decode YAML from URL: %v", err)
return
}

return
}

func readYAMLFile() (data ControlCatalog, err error) {
yamlFile, err := os.ReadFile(SourcePath)
if err != nil {
logger.Error(fmt.Sprintf("Error reading local source file: %s (%v)", SourcePath, err))
return
}

err = yaml.Unmarshal(yamlFile, &data)
if err != nil {
logger.Error("Error unmarshalling YAML file: %v", err)
return
}

return
}
2 changes: 1 addition & 1 deletion go.mod
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,9 +7,9 @@ require (
github.com/hashicorp/go-hclog v1.6.3
github.com/hashicorp/go-plugin v1.6.3
github.com/privateerproj/privateer-sdk v0.6.1
github.com/revanite-io/sci/pkg/layer2 v0.0.0-20250223164144-5fad3b3dc290
github.com/spf13/cobra v1.9.1
github.com/spf13/viper v1.19.0
gopkg.in/yaml.v2 v2.4.0
)

require (
Expand Down
3 changes: 2 additions & 1 deletion go.sum
View file
Original file line number Diff line number Diff line change
Expand Up @@ -100,6 +100,8 @@ github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 h1:Jamvg5psRI
github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
github.com/privateerproj/privateer-sdk v0.6.1 h1:uXojv/E61ETiuAU9SWprJcINA4lgV3bXK7YF/53cd68=
github.com/privateerproj/privateer-sdk v0.6.1/go.mod h1:2utlXAO2fAOFgXKLjn/t6UqkU/4MpEYy2hLxacUsIMs=
github.com/revanite-io/sci/pkg/layer2 v0.0.0-20250223164144-5fad3b3dc290 h1:F6fpBIr4ACgJm6+MOav08LFvsHbH6Xub1a4JS4VgR34=
github.com/revanite-io/sci/pkg/layer2 v0.0.0-20250223164144-5fad3b3dc290/go.mod h1:+Qqc7Q6ODGeeEqc7hNCEzJ5XtVpG+Lt6FqOamEx8MNU=
github.com/rogpeppe/go-internal v1.12.0 h1:exVL4IDcn6na9z1rAb56Vxr+CgyK3nn3O+epU5NdKM8=
github.com/rogpeppe/go-internal v1.12.0/go.mod h1:E+RYuTGaKKdloAfM02xzb0FW3Paa99yedzYV+kq4uf4=
github.com/russross/blackfriday/v2 v2.1.0/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM=
Expand Down Expand Up @@ -200,7 +202,6 @@ gopkg.in/ini.v1 v1.67.0/go.mod h1:pNLf8WUiyNEtQjuu5G5vTm06TEv9tsIgeAvK8hOrP4k=
gopkg.in/warnings.v0 v0.1.2 h1:wFXVbFY8DY5/xOe1ECiWdKCzZlxgshcYVNkBHstARME=
gopkg.in/warnings.v0 v0.1.2/go.mod h1:jksf8JmL6Qr/oQM2OXTHunEvvTAsrWBLb6OOjuVWRNI=
gopkg.in/yaml.v2 v2.2.2/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=
gopkg.in/yaml.v2 v2.4.0 h1:D8xgwECY7CYvx+Y2n4sBz93Jn9JRvxdiyyo8CTfuKaY=
gopkg.in/yaml.v2 v2.4.0/go.mod h1:RDklbk79AGWmwhnvt/jBztapEOGDOx6ZbXqjP6csGnQ=
gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
Expand Down

0 comments on commit 87f5cc7

Please sign in to comment.