diff --git a/config/config.go b/config/config.go index 6d33728..6431bb9 100644 --- a/config/config.go +++ b/config/config.go @@ -45,6 +45,20 @@ func (at *ArtifactType) UnmarshalTOML(data []byte) (err error) { return } +func (at ArtifactType) MarshalTOML() (data []byte, err error) { + s := "" + switch at { + case KernelModule: + s = "module" + case KernelExploit: + s = "exploit" + default: + err = errors.New(fmt.Sprintf("Cannot marshal %d", at)) + } + data = []byte(`"` + s + `"`) + return +} + type Artifact struct { Name string Type ArtifactType @@ -114,6 +128,22 @@ func (dt *DistroType) UnmarshalTOML(data []byte) (err error) { return } +func (dt DistroType) MarshalTOML() (data []byte, err error) { + s := "" + switch dt { + case Ubuntu: + s = "Ubuntu" + case CentOS: + s = "CentOS" + case Debian: + s = "Debian" + default: + err = errors.New(fmt.Sprintf("Cannot marshal %d", dt)) + } + data = []byte(`"` + s + `"`) + return +} + type KernelInfo struct { DistroType DistroType DistroRelease string // 18.04/7.4.1708/9.1 diff --git a/config/config_test.go b/config/config_test.go new file mode 100644 index 0000000..7bf2a59 --- /dev/null +++ b/config/config_test.go @@ -0,0 +1,30 @@ +// Copyright 2018 Mikhail Klementev. All rights reserved. +// Use of this source code is governed by a AGPLv3 license +// (or later) that can be found in the LICENSE file. + +package config + +import ( + "testing" + + "github.com/naoina/toml" +) + +func TestMarshalUnmarshal(t *testing.T) { + artifactCfg := Artifact{ + Name: "Put name here", + Type: KernelModule, + } + artifactCfg.SupportedKernels = append(artifactCfg.SupportedKernels, + KernelMask{Ubuntu, "18.04", ".*"}) + buf, err := toml.Marshal(&artifactCfg) + if err != nil { + t.Fatal(err) + } + + var artifactCfgNew Artifact + err = toml.Unmarshal(buf, &artifactCfgNew) + if err != nil { + t.Fatal(err) + } +}