From dfd1cd177a0756e37a1ed97671b81caa3c9d009f Mon Sep 17 00:00:00 2001
From: Brendan Le Glaunec <brendan.le-glaunec@epitech.eu>
Date: Sun, 5 Jul 2020 14:15:38 +0200
Subject: [PATCH] Use mapstructure instead of YAML for PortRanges

Turns out Viper does not support YAML unmarshalling directly but goes through a map structure as an intermediate: https://github.com/spf13/viper/issues/338
---
 .gitignore                         |  1 +
 pkg/configuration/configuration.go | 48 ++++++++++++++++++++----------
 pkg/configuration/portranges.go    |  2 ++
 3 files changed, 35 insertions(+), 16 deletions(-)

diff --git a/.gitignore b/.gitignore
index 66fd13c..4837d26 100644
--- a/.gitignore
+++ b/.gitignore
@@ -4,6 +4,7 @@
 *.dll
 *.so
 *.dylib
+goneypot
 
 # Test binary, built with `go test -c`
 *.test
diff --git a/pkg/configuration/configuration.go b/pkg/configuration/configuration.go
index eef9427..db7f65b 100644
--- a/pkg/configuration/configuration.go
+++ b/pkg/configuration/configuration.go
@@ -1,5 +1,7 @@
 package configuration
 
+import "errors"
+
 type Configuration struct {
 	ICMP bool       `yaml:"icmp"`
 	TCP  *TCPConfig `yaml:"tcp"`
@@ -8,19 +10,27 @@ type Configuration struct {
 }
 
 type TCPConfig struct {
-	Ports PortRanges `yaml:"ports"`
+	Ports PortRanges `mapstructure:"ports"`
 	// TODO: Fake service configuration.
 }
 
-// UnmarshalYAML implements the Unmarshaler interface of the yaml pkg.
-func (t *TCPConfig) UnmarshalYAML(unmarshal func(interface{}) error) error {
-	var yamlPortRangeSequence []string
-	if err := unmarshal(&yamlPortRangeSequence); err != nil {
-		return err
+// UnmarshalMap implements the Unmarshaler interface.
+func (t *TCPConfig) UnmarshalMap(value interface{}) error {
+	var (
+		tcpCfg map[string]interface{}
+		ok     bool
+	)
+	if tcpCfg, ok = value.(map[string]interface{}); !ok {
+		return errors.New("invalid TCP port range")
+	}
+
+	portRanges, ok := tcpCfg["ports"].([]string)
+	if !ok {
+		return errors.New("invalid TCP port range")
 	}
 
 	var err error
-	t.Ports, err = NewPortRanges(yamlPortRangeSequence)
+	t.Ports, err = NewPortRanges(portRanges)
 	if err != nil {
 		return err
 	}
@@ -29,22 +39,28 @@ func (t *TCPConfig) UnmarshalYAML(unmarshal func(interface{}) error) error {
 }
 
 type UDPConfig struct {
-	Ports PortRanges `yaml:"ports"`
+	Ports PortRanges `mapstructure:"ports"`
 	// TODO: Fake service configuration.
 }
 
-// UnmarshalYAML implements the Unmarshaler interface of the yaml pkg.
-func (u *UDPConfig) UnmarshalYAML(unmarshal func(interface{}) error) error {
-	var yamlPortRangeSequence []string
-	if err := unmarshal(&yamlPortRangeSequence); err != nil {
-		return err
+// UnmarshalMap implements the Unmarshaler interface.
+func (u *UDPConfig) UnmarshalMap(value interface{}) error {
+	var (
+		udpCfg map[string]interface{}
+		ok     bool
+	)
+	if udpCfg, ok = value.(map[string]interface{}); !ok {
+		return errors.New("invalid UDP port range")
+	}
+
+	portRanges, ok := udpCfg["ports"].([]string)
+	if !ok {
+		return errors.New("invalid UDP port range")
 	}
 
 	var err error
-	u.Ports, err = NewPortRanges(yamlPortRangeSequence)
+	t.Ports, err = NewPortRanges(portRanges)
 	if err != nil {
 		return err
 	}
-
-	return nil
 }
diff --git a/pkg/configuration/portranges.go b/pkg/configuration/portranges.go
index 4c44422..a045773 100644
--- a/pkg/configuration/portranges.go
+++ b/pkg/configuration/portranges.go
@@ -14,6 +14,8 @@ func NewPortRanges(s []string) (PortRanges, error) {
 	var result PortRanges
 
 	for _, portRange := range s {
+		fmt.Println("Port range found:", s)
+
 		ports := strings.Split(portRange, "-")
 
 		// Case where only one port is specified.