Refactored collector functionality

lkarlslund · Aug 26, 2022 · d8c68b4 · d8c68b4
1 parent febd767
commit d8c68b4
Show file tree

Hide file tree

Showing 3 changed files with 75 additions and 41 deletions.
diff --git a/collector/main.go b/collector/main.go
@@ -1,9 +1,13 @@
 package main
 
 import (
+	"encoding/json"
 	"fmt"
+	"io/ioutil"
 	"os"
+	"path/filepath"
 
+	"github.com/lkarlslund/adalanche/modules/integrations/localmachine"
 	"github.com/lkarlslund/adalanche/modules/integrations/localmachine/collect"
 	"github.com/lkarlslund/adalanche/modules/ui"
 	"github.com/lkarlslund/adalanche/modules/version"
@@ -43,15 +47,40 @@ func Execute(cmd *cobra.Command, args []string) error {
 		}
 	}
 
-	// Ensure the data folder is available
-	if _, err := os.Stat(*datapath); os.IsNotExist(err) {
-		err = os.MkdirAll(*datapath, 0711)
-		if err != nil {
-			return fmt.Errorf("Could not create data folder %v: %v", datapath, err)
-		}
+	outputpath := *datapath
+
+	err = os.MkdirAll(outputpath, 0600)
+	if err != nil {
+		return fmt.Errorf("Problem accessing output folder: %v", err)
+	}
+
+	info, err := collect.Collect()
+	if err != nil {
+		return err
+	}
+
+	if outputpath == "" {
+		ui.Warn().Msg("Missing -outputpath parameter - writing file to current directory")
+		outputpath = "."
+	}
+
+	targetname := info.Machine.Name + localmachine.Suffix
+	if info.Machine.IsDomainJoined {
+		targetname = info.Machine.Name + "$" + info.Machine.Domain + localmachine.Suffix
+	}
+	output, err := json.MarshalIndent(info, "", "  ")
+	if err != nil {
+		return fmt.Errorf("Problem marshalling JSON: %v", err)
+	}
+
+	outputfile := filepath.Join(outputpath, targetname)
+	err = ioutil.WriteFile(outputfile, output, 0600)
+	if err != nil {
+		return fmt.Errorf("Problem writing to file %v: %v", outputfile, err)
 	}
+	ui.Info().Msgf("Information collected to file %v", outputfile)
 
-	return collect.Collect(*datapath)
+	return nil
 }
 
 func main() {

diff --git a/modules/engine/securitydescriptor.go b/modules/engine/securitydescriptor.go
@@ -23,15 +23,15 @@ type ACLPermissionMask uint32
 const (
 	CONTROLFLAG_OWNER_DEFAULTED     SecurityDescriptorControlFlag = 0x0001
 	CONTROLFLAG_GROUP_DEFAULTED     SecurityDescriptorControlFlag = 0x0002
-	CONTROLFLAG_DACL_PRESENT                                      = 0x0004
-	CONTROLFLAG_DACL_DEFAULTED                                    = 0x0008
-	CONTROLFLAG_SACL_PRESENT                                      = 0x0010
-	CONTROLFLAG_SACL_DEFAULTED                                    = 0x0020
-	CONTROLFLAG_DACL_AUTO_INHERITED                               = 0x0400
-	CONTROLFLAG_SACL_AUTO_INHERITED                               = 0x0800
-	CONTROLFLAG_DACL_PROTECTED                                    = 0x1000
-	CONTROLFLAG_SACL_PROTECTED                                    = 0x2000
-	CONTROLFLAG_SELF_RELATIVE                                     = 0x8000
+	CONTROLFLAG_DACL_PRESENT        SecurityDescriptorControlFlag = 0x0004
+	CONTROLFLAG_DACL_DEFAULTED      SecurityDescriptorControlFlag = 0x0008
+	CONTROLFLAG_SACL_PRESENT        SecurityDescriptorControlFlag = 0x0010
+	CONTROLFLAG_SACL_DEFAULTED      SecurityDescriptorControlFlag = 0x0020
+	CONTROLFLAG_DACL_AUTO_INHERITED SecurityDescriptorControlFlag = 0x0400
+	CONTROLFLAG_SACL_AUTO_INHERITED SecurityDescriptorControlFlag = 0x0800
+	CONTROLFLAG_DACL_PROTECTED      SecurityDescriptorControlFlag = 0x1000
+	CONTROLFLAG_SACL_PROTECTED      SecurityDescriptorControlFlag = 0x2000
+	CONTROLFLAG_SELF_RELATIVE       SecurityDescriptorControlFlag = 0x8000
 
 	// ACE.Type
 	ACETYPE_ACCESS_ALLOWED        = 0x00

diff --git a/modules/integrations/localmachine/collect/main.go b/modules/integrations/localmachine/collect/main.go
@@ -47,15 +47,41 @@ func Execute(cmd *cobra.Command, args []string) error {
 	if op := cmd.InheritedFlags().Lookup("datapath"); op != nil {
 		outputpath = op.Value.String()
 	}
-	return Collect(outputpath)
-}
 
-func Collect(outputpath string) error {
 	err := os.MkdirAll(outputpath, 0600)
 	if err != nil {
 		return fmt.Errorf("Problem accessing output folder: %v", err)
 	}
 
+	info, err := Collect()
+	if err != nil {
+		return err
+	}
+
+	if outputpath == "" {
+		ui.Warn().Msg("Missing -outputpath parameter - writing file to current directory")
+		outputpath = "."
+	}
+
+	targetname := info.Machine.Name + localmachine.Suffix
+	if info.Machine.IsDomainJoined {
+		targetname = info.Machine.Name + "$" + info.Machine.Domain + localmachine.Suffix
+	}
+	output, err := json.MarshalIndent(info, "", "  ")
+	if err != nil {
+		return fmt.Errorf("Problem marshalling JSON: %v", err)
+	}
+
+	outputfile := filepath.Join(outputpath, targetname)
+	err = ioutil.WriteFile(outputfile, output, 0600)
+	if err != nil {
+		return fmt.Errorf("Problem writing to file %v: %v", outputfile, err)
+	}
+	ui.Info().Msgf("Information collected to file %v", outputfile)
+	return nil
+}
+
+func Collect() (localmachine.Info, error) {
 	if !is64Bit && os64Bit {
 		ui.Debug().Msgf("Running as 32-bit on 64-bit system")
 	}
@@ -745,26 +771,5 @@ func Collect(outputpath string) error {
 		Privileges: privilegesinfo,
 	}
 
-	if outputpath == "" {
-		ui.Warn().Msg("Missing -outputpath parameter - writing file to current directory")
-		outputpath = "."
-	}
-
-	targetname := info.Machine.Name + localmachine.Suffix
-	if info.Machine.IsDomainJoined {
-		targetname = info.Machine.Name + "$" + info.Machine.Domain + localmachine.Suffix
-	}
-	output, err := json.MarshalIndent(info, "", "  ")
-	if err != nil {
-		return fmt.Errorf("Problem marshalling JSON: %v", err)
-	}
-
-	outputfile := filepath.Join(outputpath, targetname)
-	err = ioutil.WriteFile(outputfile, output, 0600)
-	if err != nil {
-		return fmt.Errorf("Problem writing to file %v: %v", outputfile, err)
-	}
-	ui.Info().Msgf("Information collected to file %v", outputfile)
-
-	return nil
+	return info, nil
 }