Skip to content

Commit

Permalink
refactor: divide up large library files into single responsibility fi…
Browse files Browse the repository at this point in the history
…les (#53)

This commit moves the previous functions in agentClient and clientUtils into purpose-specific files. The overall codebase is now more maintainable and easier to navigate.

Signed-off-by: UncleSp1d3r <unclespider@protonmail.com>
  • Loading branch information
unclesp1d3r authored Oct 4, 2024
1 parent a1f6257 commit bbffc89
Show file tree
Hide file tree
Showing 9 changed files with 1,069 additions and 984 deletions.
764 changes: 21 additions & 743 deletions lib/agentClient.go

Large diffs are not rendered by default.

177 changes: 177 additions & 0 deletions lib/benchmarkManager.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,177 @@
package lib

import (
"errors"
"fmt"
"github.com/duke-git/lancet/v2/convertor"
"github.com/duke-git/lancet/v2/strutil"
"github.com/unclesp1d3r/cipherswarm-agent-sdk-go/models/components"
"github.com/unclesp1d3r/cipherswarm-agent-sdk-go/models/operations"
"github.com/unclesp1d3r/cipherswarmagent/lib/arch"
"github.com/unclesp1d3r/cipherswarmagent/lib/hashcat"
"github.com/unclesp1d3r/cipherswarmagent/shared"
"net/http"
"strings"
)

// sendBenchmarkResults sends the collected benchmark results to a server endpoint.
// It converts each benchmarkResult into a HashcatBenchmark and appends them to a slice.
// If the conversion fails for a result, it continues to the next result.
// Creates a SubmitBenchmarkRequestBody with the HashcatBenchmarks slice and submits it via SdkClient.
// Returns an error if submission or the response received is not successful.
func sendBenchmarkResults(benchmarkResults []benchmarkResult) error {
var benchmarks []components.HashcatBenchmark //nolint:prealloc

for _, result := range benchmarkResults {
benchmark, err := createBenchmark(result)
if err != nil {
continue
}
benchmarks = append(benchmarks, benchmark)
}

results := operations.SubmitBenchmarkRequestBody{
HashcatBenchmarks: benchmarks,
}

res, err := SdkClient.Agents.SubmitBenchmark(Context, shared.State.AgentID, results)
if err != nil {
return err
}

if res.StatusCode == http.StatusNoContent {
return nil
}

return errors.New("bad response: " + res.RawResponse.Status)
}

// createBenchmark converts a benchmarkResult to a components.HashcatBenchmark struct.
// It handles the conversion of string fields in benchmarkResult to appropriate types.
// Returns a HashcatBenchmark instance and an error if any conversion fails.
func createBenchmark(result benchmarkResult) (components.HashcatBenchmark, error) {
hashType, err := convertor.ToInt(result.HashType)
if err != nil {
return components.HashcatBenchmark{}, fmt.Errorf("failed to convert HashType: %w", err)
}
runtimeMs, err := convertor.ToInt(result.RuntimeMs)
if err != nil {
return components.HashcatBenchmark{}, fmt.Errorf("failed to convert RuntimeMs: %w", err)
}
speedHs, err := convertor.ToFloat(result.SpeedHs)
if err != nil {
return components.HashcatBenchmark{}, fmt.Errorf("failed to convert SpeedHs: %w", err)
}
device, err := convertor.ToInt(result.Device)
if err != nil {
return components.HashcatBenchmark{}, fmt.Errorf("failed to convert Device: %w", err)
}

return components.HashcatBenchmark{
HashType: hashType,
Runtime: runtimeMs,
HashSpeed: speedHs,
Device: device,
}, nil
}

// UpdateBenchmarks updates the benchmark metrics using Hashcat.
// Creates a Hashcat session with benchmark parameters and initiates the benchmarking process.
// Logs the session start, runs the benchmark task, and updates the results.
// If any errors occur during session creation or result sending, logs the errors and returns them.
func UpdateBenchmarks() error {
jobParams := hashcat.Params{
AttackMode: hashcat.AttackBenchmark,
AdditionalArgs: arch.GetAdditionalHashcatArgs(),
BackendDevices: Configuration.Config.BackendDevices,
OpenCLDevices: Configuration.Config.OpenCLDevices,
EnableAdditionalHashTypes: shared.State.EnableAdditionalHashTypes,
}

sess, err := hashcat.NewHashcatSession("benchmark", jobParams)
if err != nil {
return logAndSendError("Failed to create benchmark session", err, operations.SeverityMajor, nil)
}
shared.Logger.Debug("Starting benchmark session", "cmdline", sess.CmdLine())

displayBenchmarkStarting()
benchmarkResult, done := runBenchmarkTask(sess)
if done {
return nil
}
displayBenchmarksComplete(benchmarkResult)
if err := sendBenchmarkResults(benchmarkResult); err != nil {
return logAndSendError("Error updating benchmarks", err, operations.SeverityCritical, nil)
}

return nil
}

// runBenchmarkTask starts a hashcat benchmark session and processes its output.
// It returns a slice of benchmark results and a boolean indicating an error state.
func runBenchmarkTask(sess *hashcat.Session) ([]benchmarkResult, bool) {
err := sess.Start()
if err != nil {
shared.Logger.Error("Failed to start benchmark session", "error", err)

return nil, true
}

var benchmarkResults []benchmarkResult
waitChan := make(chan int)

go func() {
defer close(waitChan)
for {
select {
case stdOutLine := <-sess.StdoutLines:
handleBenchmarkStdOutLine(stdOutLine, &benchmarkResults)
case stdErrLine := <-sess.StderrMessages:
handleBenchmarkStdErrLine(stdErrLine)
case statusUpdate := <-sess.StatusUpdates:
shared.Logger.Debug("Benchmark status update", "status", statusUpdate) // This should never happen
case crackedHash := <-sess.CrackedHashes:
shared.Logger.Debug("Benchmark cracked hash", "hash", crackedHash) // This should never happen
case err := <-sess.DoneChan:
if err != nil {
shared.Logger.Error("Benchmark session failed", "error", err)
SendAgentError(err.Error(), nil, operations.SeverityFatal)
}

return
}
}
}()

<-waitChan

return benchmarkResults, false
}

// handleBenchmarkStdOutLine processes a line of benchmark output, extracting relevant data and appending it to result.
func handleBenchmarkStdOutLine(line string, results *[]benchmarkResult) {
fields := strings.Split(line, ":")
if len(fields) != 6 {
shared.Logger.Debug("Unknown benchmark line", "line", line)

return
}

result := benchmarkResult{
Device: fields[0],
HashType: fields[1],
RuntimeMs: fields[3],
HashTimeMs: fields[4],
SpeedHs: fields[5],
}
displayBenchmark(result)
*results = append(*results, result)
}

// handleBenchmarkStdErrLine processes each line from the benchmark's standard error output, logs it, and reports warnings to the server.
func handleBenchmarkStdErrLine(line string) {
displayBenchmarkError(line)
if strutil.IsNotBlank(line) {
SendAgentError(line, nil, operations.SeverityWarning)
}
}
101 changes: 0 additions & 101 deletions lib/clientUtils.go
Original file line number Diff line number Diff line change
@@ -1,7 +1,6 @@
package lib

import (
"context"
"encoding/hex"
"errors"
"io"
Expand All @@ -16,14 +15,11 @@ import (
"github.com/duke-git/lancet/v2/cryptor"
"github.com/duke-git/lancet/v2/fileutil"
"github.com/duke-git/lancet/v2/strutil"
"github.com/duke-git/lancet/v2/validator"
"github.com/hashicorp/go-getter"
"github.com/shirou/gopsutil/v3/process"
"github.com/spf13/viper"
"github.com/unclesp1d3r/cipherswarm-agent-sdk-go/models/components"
"github.com/unclesp1d3r/cipherswarm-agent-sdk-go/models/operations"
"github.com/unclesp1d3r/cipherswarmagent/lib/arch"
"github.com/unclesp1d3r/cipherswarmagent/lib/utils"
"github.com/unclesp1d3r/cipherswarmagent/shared"
)

Expand Down Expand Up @@ -259,103 +255,6 @@ func writeResponseToFile(responseStream io.Reader, filePath string) error {
return nil
}

// downloadFile downloads a file from a given URL and saves it to the specified path with optional checksum verification.
// If the URL is invalid, it returns an error. If the file already exists and the checksum matches, the download is skipped.
func downloadFile(url string, path string, checksum string) error {
if !validator.IsUrl(url) {
shared.Logger.Error("Invalid URL", "url", url)

return errors.New("invalid URL")
}

if fileExistsAndValid(path, checksum) {
shared.Logger.Info("Download already exists", "path", path)

return nil
}

displayDownloadFile(url, path)
if err := downloadAndVerifyFile(url, path, checksum); err != nil {
return err
}
displayDownloadFileComplete(url, path)

return nil
}

// fileExistsAndValid checks if a file exists at the given path and, if a checksum is provided, verifies its validity.
// The function returns true if the file exists and matches the given checksum, or if no checksum is provided.
// If the file does not exist or the checksum verification fails, appropriate error messages are logged.
func fileExistsAndValid(path string, checksum string) bool {
if !fileutil.IsExist(path) {
return false
}

if strutil.IsBlank(checksum) {
return true
}

fileChecksum, err := cryptor.Md5File(path)
if err != nil {
shared.Logger.Error("Error calculating file checksum", "path", path, "error", err)

return false
}

if fileChecksum == checksum {
return true
}

shared.Logger.Warn("Checksums do not match", "path", path, "url_checksum", checksum, "file_checksum", fileChecksum)
SendAgentError("Resource "+path+" exists, but checksums do not match", nil, operations.SeverityInfo)
if err := os.Remove(path); err != nil {
SendAgentError(err.Error(), nil, operations.SeverityMajor)
shared.Logger.Error("Error removing file with mismatched checksum", "path", path, "error", err)
}

return false
}

// downloadAndVerifyFile downloads a file from the given URL and saves it to the specified path, verifying the checksum if provided.
// If a checksum is given, it is appended to the URL before download. The function then configures a client for secure file transfer.
// The file is downloaded using the configured client. After downloading, the file's checksum is verified, if provided, to ensure integrity.
// If the checksum does not match, an error is returned, indicating the downloaded file is corrupt.
func downloadAndVerifyFile(url string, path string, checksum string) error {
if strutil.IsNotBlank(checksum) {
var err error
url, err = appendChecksumToURL(url, checksum)
if err != nil {
return err
}
}

client := &getter.Client{
Ctx: context.Background(),
Dst: path,
Src: url,
Pwd: shared.State.CrackersPath,
Insecure: true,
Mode: getter.ClientModeFile,
}

_ = client.Configure(
getter.WithProgress(utils.DefaultProgressBar),
getter.WithUmask(os.FileMode(0o022)),
)

if err := client.Get(); err != nil {
shared.Logger.Debug("Error downloading file", "error", err)

return err
}

if strutil.IsNotBlank(checksum) && !fileExistsAndValid(path, checksum) {
return errors.New("downloaded file checksum does not match")
}

return nil
}

// appendChecksumToURL appends the given checksum as a query parameter to the provided URL and returns the updated URL.
// If there is an error parsing the URL, it logs the error and returns an empty string and the error.
func appendChecksumToURL(url string, checksum string) (string, error) {
Expand Down
82 changes: 82 additions & 0 deletions lib/crackerUtils.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,82 @@
package lib

import (
"github.com/duke-git/lancet/v2/fileutil"
"github.com/spf13/viper"
"github.com/unclesp1d3r/cipherswarm-agent-sdk-go/models/operations"
"github.com/unclesp1d3r/cipherswarmagent/shared"
"net/http"
"path"
)

// setNativeHashcatPath sets the path for the native Hashcat binary if it is found in the system, otherwise logs and reports error.
func setNativeHashcatPath() error {
shared.Logger.Debug("Using native Hashcat")
binPath, err := findHashcatBinary()
if err != nil {
shared.Logger.Error("Error finding hashcat binary: ", err)
SendAgentError(err.Error(), nil, operations.SeverityCritical)

return err
}
shared.Logger.Info("Found Hashcat binary", "path", binPath)
viper.Set("hashcat_path", binPath)

return viper.WriteConfig()
}

// UpdateCracker checks for updates to the cracker and applies them if available.
// It starts by logging the beginning of the update process and attempts to fetch the current version of Hashcat.
// It then calls the API to check if there are any updates available. Depending on the API response, it either handles
// the update process or logs the absence of any new updates. If any errors occur during these steps, they are logged and handled accordingly.
func UpdateCracker() {
shared.Logger.Info("Checking for updated cracker")
currentVersion, err := getCurrentHashcatVersion()
if err != nil {
shared.Logger.Error("Error getting current hashcat version", "error", err)

return
}

response, err := SdkClient.Crackers.CheckForCrackerUpdate(Context, &agentPlatform, &currentVersion)
if err != nil {
handleAPIError("Error connecting to the CipherSwarm API", err, operations.SeverityCritical)

return
}

if response.StatusCode == http.StatusNoContent {
shared.Logger.Debug("No new cracker available")

return
}

if response.StatusCode == http.StatusOK {
update := response.GetCrackerUpdate()
if update.GetAvailable() {
_ = handleCrackerUpdate(update)
} else {
shared.Logger.Debug("No new cracker available", "latest_version", update.GetLatestVersion())
}
} else {
shared.Logger.Error("Error checking for updated cracker", "CrackerUpdate", response.RawResponse.Status)
}
}

// validateHashcatDirectory checks if the given hashcat directory exists and contains the specified executable.
func validateHashcatDirectory(hashcatDirectory, execName string) bool {
if !fileutil.IsDir(hashcatDirectory) {
shared.Logger.Error("New hashcat directory does not exist", "path", hashcatDirectory)

return false
}

hashcatBinaryPath := path.Join(hashcatDirectory, execName)
if !fileutil.IsExist(hashcatBinaryPath) {
shared.Logger.Error("New hashcat binary does not exist", "path", hashcatBinaryPath)

return false
}

return true
}
Loading

0 comments on commit bbffc89

Please sign in to comment.