Skip to content

Commit

Permalink
Supports JSON for rules import/export (#5)
Browse files Browse the repository at this point in the history
* wip

* Supports JSON
  • Loading branch information
Ryxias committed Sep 9, 2021
1 parent 8f5f40a commit 1ebd59b
Show file tree
Hide file tree
Showing 6 changed files with 249 additions and 54 deletions.
16 changes: 16 additions & 0 deletions examples/sample-rules2.json
@@ -0,0 +1,16 @@
[
{
"sha256": "d84db96af8c2e60ac4c851a21ec460f6f84e0235beb17d24a78712b9b021ed57",
"type": "CERTIFICATE",
"policy": "ALLOWLIST",
"custom_msg": "",
"description": "Software Signing by Apple Inc."
},
{
"sha256": "345a8e098bd04794aaeefda8c9ef56a0bf3d3706d67d35bc0e23f11bb3bffce5",
"type": "CERTIFICATE",
"policy": "ALLOWLIST",
"custom_msg": "",
"description": "Developer ID Application: Google, Inc. (EQHXZ8M8AV)"
}
]
122 changes: 106 additions & 16 deletions internal/cli/rules/rules-export.go
@@ -1,16 +1,22 @@
package rules

import (
"encoding/json"
"fmt"
"io/ioutil"

"github.com/pkg/errors"
"github.com/spf13/cobra"

"github.com/airbnb/rudolph/internal/csv"
"github.com/airbnb/rudolph/pkg/dynamodb"
"github.com/airbnb/rudolph/pkg/model/globalrules"
"github.com/airbnb/rudolph/pkg/types"
)

func addRuleExportCommand() {
var filename string
var format string
var ruleExportCmd = &cobra.Command{
Use: "export <file-name>",
Aliases: []string{"rules-export"},
Expand All @@ -23,19 +29,74 @@ func addRuleExportCommand() {

dynamodbClient := dynamodb.GetClient(table, region)

return runExport(dynamodbClient, filename)
return runExport(dynamodbClient, filename, format)
},
}

ruleExportCmd.Flags().StringVarP(&filename, "filename", "f", "", "The filename")
_ = ruleExportCmd.MarkFlagRequired("filename")

ruleExportCmd.Flags().StringVarP(&format, "fileformat", "t", "csv", "File format (one of: [json|csv])")

RulesCmd.AddCommand(ruleExportCmd)
}

func runExport(
client dynamodb.QueryAPI,
filename string,
format string,
) (err error) {
switch format {
case "json":
return runJsonExport(client, filename)
case "csv":
return runCsvExport(client, filename)
}
return
}

type fileRule struct {
RuleType types.RuleType `json:"type"`
Policy types.Policy `json:"policy"`
SHA256 string `json:"sha256"`
CustomMessage string `json:"custom_msg,omitempty"`
Description string `json:"description"`
}

func runJsonExport(client dynamodb.QueryAPI, filename string) (err error) {
var jsonRules []fileRule
fmt.Println("Querying rules from DynamoDB...")
total, err := getRules(client, func(rule globalrules.GlobalRuleRow) (err error) {
jsonRules = append(jsonRules, fileRule{
SHA256: rule.SHA256,
RuleType: rule.RuleType,
Policy: rule.Policy,
CustomMessage: rule.CustomMessage,
Description: rule.Description,
})
return
})
if err != nil {
return
}

jsondata, err := json.MarshalIndent(jsonRules, "", " ")
if err != nil {
return
}
err = ioutil.WriteFile(filename, jsondata, 0644)
if err != nil {
return
}

fmt.Printf("rules discovered: %d, rules written: %d\n", total, len(jsonRules))

return
}

func runCsvExport(
client dynamodb.QueryAPI,
filename string,
) (err error) {
csvRules := make(chan []string)

Expand All @@ -52,6 +113,45 @@ func runExport(
panic(err)
}

fmt.Println("Querying rules from DynamoDB...")
var totalWritten int64
total, err := getRules(client, func(rule globalrules.GlobalRuleRow) (err error) {
ruleType, err := rule.RuleType.MarshalText()
if err != nil {
return
}
policy, err := rule.Policy.MarshalText()
if err != nil {
return
}
record := []string{
rule.SHA256,
string(ruleType),
string(policy),
rule.CustomMessage,
rule.Description,
}
if err != nil {
return
}

totalWritten += 1
csvRules <- record
return
})
if err != nil {
return
}

close(csvRules)
wg.Wait()

fmt.Printf("rules discovered: %d, rules written: %d\n", total, totalWritten)

return
}

func getRules(client dynamodb.QueryAPI, callback func(globalrules.GlobalRuleRow) error) (total int64, err error) {
var key *dynamodb.PrimaryKey
for {
rules, nextkey, inerr := globalrules.GetPaginatedGlobalRules(client, 50, key)
Expand All @@ -64,27 +164,17 @@ func runExport(
}

for _, rule := range *rules {
ruleType, _ := rule.RuleType.MarshalText()
policy, _ := rule.Policy.MarshalText()
record := []string{
rule.SHA256,
string(ruleType),
string(policy),
rule.CustomMessage,
rule.Description,
total += 1
err = callback(rule)
if err != nil {
return
}

csvRules <- record
}

if nextkey == nil {
break
}
key = nextkey
}
close(csvRules)

wg.Wait()

return nil
return
}

0 comments on commit 1ebd59b

Please sign in to comment.