forked from Velocidex/velociraptor
-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Refactor config parsing into a builder pattern. (Velocidex#381)
Makes it clear how we fetch the config from different commands. Depending on the command, some commands accept api config, some rely on embedded config etc. This refactor makes it much clearer which config each command support and how to provide it. Also added integration tests for repacking, autoexec and config loading order.
- Loading branch information
Showing
53 changed files
with
1,100 additions
and
648 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file was deleted.
Oops, something went wrong.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,232 @@ | ||
| // Tests for the binary. | ||
|
|
||
| package main | ||
|
|
||
| import ( | ||
| "bytes" | ||
| "fmt" | ||
| "io/ioutil" | ||
| "os" | ||
| "os/exec" | ||
| "path/filepath" | ||
| "runtime" | ||
| "testing" | ||
|
|
||
| "github.com/alecthomas/assert" | ||
| "github.com/stretchr/testify/require" | ||
| "github.com/stretchr/testify/suite" | ||
| "github.com/tink-ab/tempfile" | ||
| "www.velocidex.com/golang/velociraptor/config" | ||
| "www.velocidex.com/golang/velociraptor/constants" | ||
| ) | ||
|
|
||
| type MainTestSuite struct { | ||
| suite.Suite | ||
| binary string | ||
| extension string | ||
| } | ||
|
|
||
| func (self *MainTestSuite) SetupTest() { | ||
| if runtime.GOOS == "windows" { | ||
| self.extension = ".exe" | ||
| } | ||
|
|
||
| // Search for a valid binary to run. | ||
| binaries, err := filepath.Glob( | ||
| "../output/velociraptor*" + constants.VERSION + "-" + runtime.GOOS + | ||
| "-" + runtime.GOARCH + self.extension) | ||
| assert.NoError(self.T(), err) | ||
|
|
||
| if len(binaries) == 0 { | ||
| binaries, _ = filepath.Glob("../output/velociraptor*" + | ||
| self.extension) | ||
| } | ||
|
|
||
| self.binary = binaries[0] | ||
| fmt.Printf("Found binary %v\n", self.binary) | ||
| } | ||
|
|
||
| func (self *MainTestSuite) TestHelp() { | ||
| cmd := exec.Command(self.binary, "--help") | ||
| out, err := cmd.CombinedOutput() | ||
| require.NoError(self.T(), err) | ||
| require.Contains(self.T(), string(out), "usage:") | ||
| } | ||
|
|
||
| const autoexec_file = ` | ||
| autoexec: | ||
| argv: | ||
| - artifacts | ||
| - list | ||
| artifact_definitions: | ||
| - name: MySpecialArtifact | ||
| sources: | ||
| - query: SELECT * FROM info() | ||
| - name: Windows.Sys.Interfaces | ||
| description: Override a built in artifact in the config. | ||
| sources: | ||
| - query: SELECT "MySpecialInterface" FROM scope() | ||
| ` | ||
|
|
||
| func (self *MainTestSuite) TestAutoexec() { | ||
| // Create a tempfile for the repacked binary. | ||
| exe, err := tempfile.TempFile("", "exe", self.extension) | ||
| assert.NoError(self.T(), err) | ||
|
|
||
| defer os.Remove(exe.Name()) | ||
| exe.Close() | ||
|
|
||
| // A temp file for the config. | ||
| config_file, err := ioutil.TempFile("", "config") | ||
| assert.NoError(self.T(), err) | ||
|
|
||
| defer os.Remove(config_file.Name()) | ||
| config_file.Write([]byte(autoexec_file)) | ||
| config_file.Close() | ||
|
|
||
| // Repack the config in the binary. | ||
| cmd := exec.Command(self.binary, "config", "repack", config_file.Name(), exe.Name()) | ||
| out, err := cmd.Output() | ||
| require.NoError(self.T(), err) | ||
|
|
||
| // Run the repacked binary with no args - it should run the | ||
| // `artifacts list` command. | ||
| cmd = exec.Command(exe.Name()) | ||
| out, err = cmd.Output() | ||
| require.NoError(self.T(), err) | ||
|
|
||
| // The output should contain MySpecialArtifact as well as the | ||
| // standard artifacts. | ||
| require.Contains(self.T(), string(out), "MySpecialArtifact") | ||
| require.Contains(self.T(), string(out), "Windows.Sys.Users") | ||
|
|
||
| // If provided args it works normally. | ||
| cmd = exec.Command(exe.Name(), | ||
| "artifacts", "collect", "Windows.Sys.Interfaces", "--format", "jsonl") | ||
| out, err = cmd.Output() | ||
| require.NoError(self.T(), err) | ||
|
|
||
| // Config artifacts override built in artifacts. | ||
| require.Contains(self.T(), string(out), "MySpecialInterface") | ||
| } | ||
|
|
||
| func (self *MainTestSuite) TestGenerateConfigWithMerge() { | ||
| // A temp file for the generated config. | ||
| config_file, err := ioutil.TempFile("", "config") | ||
| assert.NoError(self.T(), err) | ||
|
|
||
| defer os.Remove(config_file.Name()) | ||
| defer config_file.Close() | ||
|
|
||
| cmd := exec.Command( | ||
| self.binary, "config", "generate", "--merge", | ||
| `{"Client": {"nonce": "Foo", "writeback_linux": "some_location"}}`) | ||
| out, err := cmd.Output() | ||
| require.NoError(self.T(), err) | ||
|
|
||
| // Write the config to the tmp file | ||
| config_file_content := out | ||
| config_file.Write(out) | ||
| config_file.Close() | ||
|
|
||
| // Try to load it now. | ||
| config_obj, err := new(config.Loader).WithFileLoader(config_file.Name()). | ||
| WithRequiredFrontend().WithRequiredCA().WithRequiredClient(). | ||
| LoadAndValidate() | ||
| require.NoError(self.T(), err) | ||
|
|
||
| require.Equal(self.T(), config_obj.Client.Nonce, "Foo") | ||
| require.Equal(self.T(), config_obj.Client.WritebackLinux, "some_location") | ||
|
|
||
| // Try to show a config without a flag. | ||
| cmd = exec.Command(self.binary, "config", "show") | ||
| cmd.Env = append(os.Environ(), | ||
| "VELOCIRAPTOR_CONFIG=", | ||
| ) | ||
| out, err = cmd.Output() | ||
| require.Error(self.T(), err) | ||
|
|
||
| // Specify the config on the commandline | ||
| cmd = exec.Command(self.binary, "config", "show", "--config", config_file.Name()) | ||
| cmd.Env = append(os.Environ(), | ||
| "VELOCIRAPTOR_CONFIG=", | ||
| ) | ||
| out, err = cmd.Output() | ||
| require.NoError(self.T(), err) | ||
| require.Contains(self.T(), string(out), "Foo") | ||
|
|
||
| // Specify the config in the environment | ||
| cmd = exec.Command(self.binary, "config", "show") | ||
| cmd.Env = append(os.Environ(), | ||
| "VELOCIRAPTOR_CONFIG="+config_file.Name(), | ||
| ) | ||
| out, err = cmd.Output() | ||
| require.NoError(self.T(), err) | ||
| require.Contains(self.T(), string(out), "Foo") | ||
|
|
||
| // Specifying invalid config in the flag is a hard stop - even | ||
| // if there is a valid environ. | ||
| cmd = exec.Command(self.binary, "config", "show", "--config", "XXXX") | ||
| cmd.Env = append(os.Environ(), | ||
| "VELOCIRAPTOR_CONFIG="+config_file.Name(), | ||
| ) | ||
| out, err = cmd.Output() | ||
| require.Error(self.T(), err) | ||
|
|
||
| // Create a tempfile for the repacked binary. | ||
| exe, err := tempfile.TempFile("", "exe", self.extension) | ||
| assert.NoError(self.T(), err) | ||
|
|
||
| defer os.Remove(exe.Name()) | ||
| exe.Close() | ||
|
|
||
| // Repack the config in the binary. | ||
| cmd = exec.Command(self.binary, "config", "repack", config_file.Name(), exe.Name()) | ||
| out, err = cmd.Output() | ||
| require.NoError(self.T(), err) | ||
|
|
||
| // Run the repacked binary with invalid environ - config | ||
| // should come from embedded. | ||
| cmd = exec.Command(exe.Name(), "config", "show") | ||
| cmd.Env = append(os.Environ(), | ||
| "VELOCIRAPTOR_CONFIG=XXXX", | ||
| ) | ||
| out, err = cmd.Output() | ||
| require.NoError(self.T(), err) | ||
| require.Contains(self.T(), string(out), "Foo") | ||
|
|
||
| // Make second copy of config file and store modified version | ||
| second_config_file, err := ioutil.TempFile("", "config") | ||
| assert.NoError(self.T(), err) | ||
|
|
||
| defer os.Remove(second_config_file.Name()) | ||
|
|
||
| // Second file has no Foo in it | ||
| second_config_file_content := bytes.ReplaceAll( | ||
| config_file_content, []byte(`Foo`), []byte(`Bar`)) | ||
| second_config_file.Write(second_config_file_content) | ||
| second_config_file.Close() | ||
|
|
||
| // Check that embedded binary with config will use its | ||
| // embedded version, even if an env var is specified. | ||
| cmd = exec.Command(exe.Name(), "config", "show") | ||
| cmd.Env = append(os.Environ(), | ||
| "VELOCIRAPTOR_CONFIG="+second_config_file.Name(), | ||
| ) | ||
| out, err = cmd.Output() | ||
| require.NoError(self.T(), err) | ||
| require.Contains(self.T(), string(out), "Foo") | ||
|
|
||
| // If a config is explicitly specified, it will override even the | ||
| // embedded config. | ||
| cmd = exec.Command(exe.Name(), "config", "show", "--config", second_config_file.Name()) | ||
| out, err = cmd.Output() | ||
| require.NoError(self.T(), err) | ||
|
|
||
| // loaded config contains Bar | ||
| require.Contains(self.T(), string(out), "Bar") | ||
| } | ||
|
|
||
| func TestMainProgram(t *testing.T) { | ||
| suite.Run(t, &MainTestSuite{}) | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.