Merge pull request #3 from ohsawa0515/improve_config

Improve rewrite in config file.

README.md

@@ -81,6 +81,56 @@ host = "192.0.2.11"
 ### ec2-vuls-config end ###
 ```
 
+### Tags
+
+It can be reflected in config by setting a tag such as `vuls:user`, `vuls:port` and so on.
+
+`<...>` is the name of tag.
+
+```toml
+[servers]
+
+[servers.<Name>]
+host = "<<Private IP address of instance>>"
+port = "<vuls:port>"
+user = "<vuls:user>"
+keyPath = "<vuls:keyPath>"
+
+# Set value of tag as comma-separated.
+cpeNames = [
+<vuls:cpeNames>
+]
+
+# Set value of tag as comma-separated.
+ignoreCves = [
+<vuls:ignoreCves>
+]
+
+# Example
+
+`vuls:user` => vuls
+`vuls:port` => 22
+`vuls:keyPath` => /opt/vuls/.ssh/id_rsa
+`vuls:cpeNames` => cpe:/a:rubyonrails:ruby_on_rails:4.2.7.1,cpe:/a:rubyonrails:ruby_on_rails:4.2.8,cpe:/a:rubyonrails:ruby_on_rails:5.0.1
+`vuls:ignoreCves` => CVE-2014-2913,CVE-2016-6314
+
+[servers.web-server-1]
+host = "192.0.2.11"
+user = "vuls"
+port = "22"
+keyPath = "/opt/vuls/.ssh/id_rsa"
+cpeNames = [
+"cpe:/a:rubyonrails:ruby_on_rails:4.2.7.1",
+"cpe:/a:rubyonrails:ruby_on_rails:4.2.8",
+"cpe:/a:rubyonrails:ruby_on_rails:5.0.1",
+]
+ignoreCves = [
+"CVE-2014-2913",
+"CVE-2016-6314",
+]
+```
+
+
 ### Options
 
 #### --config (-c)

contents.go

@@ -4,6 +4,7 @@ import (
 	"io/ioutil"
 	"os"
 	"regexp"
+	"strings"
 	"time"
 
 	"github.com/aws/aws-sdk-go/service/ec2"
@@ -16,29 +17,70 @@ const (
 
 func GenerateServerSection(instances []*ec2.Instance) []byte {
 	b := make([]byte, 0, 1024)
+	b = append(b, START+"\n"...)
+	b = append(b, "# Updated "+time.Now().Format(time.RFC3339)+"\n\n"...)
 	for _, instance := range instances {
-		b = append(b, "[servers."+GetTagValue(instance, "Name")+"]\n"...)
+
+		if name := GetTagValue(instance, "Name"); name != nil {
+			b = append(b, "[servers."+*name+"]\n"...)
+		} else {
+			continue
+		}
 		b = append(b, "host = \""+*instance.PrivateIpAddress+"\"\n"...)
+
+		if port := GetTagValue(instance, "vuls:port"); port != nil {
+			b = append(b, "port = \""+*port+"\"\n"...)
+		}
+
+		if user := GetTagValue(instance, "vuls:user"); user != nil {
+			b = append(b, "user = \""+*user+"\"\n"...)
+		}
+
+		if keyPath := GetTagValue(instance, "vuls:keyPath"); keyPath != nil {
+			b = append(b, "keyPath = \""+*keyPath+"\"\n"...)
+		}
+
+		if cpeNames := GetTagValue(instance, "vuls:cpeNames"); cpeNames != nil {
+			b = append(b, "cpeNames = [\n"...)
+			for _, cpeName := range strings.Split(*cpeNames, ",") {
+				b = append(b, "\""+cpeName+"\",\n"...)
+			}
+			b = append(b, "]\n"...)
+		}
+
+		if ignoreCves := GetTagValue(instance, "vuls:ignoreCves"); ignoreCves != nil {
+			b = append(b, "ignoreCves = [\n"...)
+			for _, ignoreCve := range strings.Split(*ignoreCves, ",") {
+				b = append(b, "\""+ignoreCve+"\",\n"...)
+			}
+			b = append(b, "]\n"...)
+		}
+
 		b = append(b, "\n"...)
 	}
+	b = append(b, END...)
 	return b
 }
 
-func CreateConfig(content []byte, config []byte) []byte {
-	re := regexp.MustCompile("(?m)" + START + "[\\s\\S]*?" + END)
+func MergeConfig(currentConfig, newConfig []byte) []byte {
 
-	b := make([]byte, 0, 1024)
-	b = append(b, START+"\n"...)
-	b = append(b, "# Updated "+time.Now().Format(time.RFC3339)+"\n\n"...)
-	b = append(b, content...)
-	b = append(b, END...)
+	// If it has already been created, it is rewritten.
+	re := regexp.MustCompile("(?m)" + START + "[\\s\\S]*?" + END)
+	if re.Match(currentConfig) {
+		return re.ReplaceAll(currentConfig, newConfig)
+	}
 
-	// if match, return replaced contents
-	if re.Match(config) {
-		return re.ReplaceAll(config, b)
+	// If it finds servers section, it is appended.
+	re = regexp.MustCompile("(?m)\\[servers.*\\][\\s\\S]*")
+	if re.Match(currentConfig) {
+		currentConfig = append(currentConfig, newConfig...)
+		return currentConfig
 	}
-	config = append(config, b...)
-	return config
+
+	// In the case that it doesn't finds servers section.
+	currentConfig = append(currentConfig, []byte("[servers]\n")...)
+	currentConfig = append(currentConfig, newConfig...)
+	return currentConfig
 }
 
 func LoadFile(path string) ([]byte, error) {

ec2-vuls-config.go

@@ -48,7 +48,7 @@ func main() {
 			return cli.NewExitError(err.Error(), 1)
 		}
 
-		new_config := CreateConfig(GenerateServerSection(instances), config)
+		new_config := MergeConfig(config, GenerateServerSection(instances))
 		if c.Bool("print") {
 			fmt.Println(string(new_config))
 		} else {

ec2.go

@@ -92,11 +92,11 @@ func DescribeInstances(filters string) ([]*ec2.Instance, error) {
 }
 
 // GetTagValue returns value of EC2 tag.
-func GetTagValue(instance *ec2.Instance, tag_name string) string {
+func GetTagValue(instance *ec2.Instance, tagName string) *string {
 	for _, t := range instance.Tags {
-		if *t.Key == tag_name {
-			return *t.Value
+		if *t.Key == tagName {
+			return t.Value
 		}
 	}
-	return ""
+	return nil
 }