Skip to content

Commit

Permalink
Moved agent helper functions to separate file
Browse files Browse the repository at this point in the history
  • Loading branch information
VikingPingvin committed Dec 4, 2020
1 parent a9033f8 commit c87ae22
Showing 1 changed file with 210 additions and 0 deletions.
210 changes: 210 additions & 0 deletions locker/agent_functions.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,210 @@
package locker

import (
"bufio"
"crypto/sha256"
"encoding/binary"
"errors"
"fmt"
"io"
"net"
"os"
"path/filepath"
"strings"
"vikingPingvin/locker/locker/messaging"
"vikingPingvin/locker/locker/messaging/protobuf"

"github.com/rs/xid"
"github.com/rs/zerolog/log"
"google.golang.org/protobuf/proto"
)

// If --file cli input is not null, parse file
func parseAndSendMetaData(connection net.Conn, inputData *InputData) (fileInfo os.FileInfo, err error) {
fileInfo, err = os.Stat(inputData.FilePath)
if os.IsNotExist(err) {
log.Error().Msgf("Parsing file Input error: %v", err)
return fileInfo, err
}
inputData.FileHash = hashFile(inputData.FilePath)
inputData.FileName = fileInfo.Name()

log.Info().
Str("file name", fileInfo.Name()).
Str("Namespace", fmt.Sprintf("%s/%s/%s", inputData.NameSpace, inputData.Project, inputData.JobID)).
Str("size", fmt.Sprintf("%d", fileInfo.Size())).
Str("hash", fmt.Sprintf("%v", inputData.FileHash)).
Str("id", inputData.ID.String()).
Msg("Artifact metadata parsing finished")

message, err := messaging.CreateMessage_FileMeta(
inputData.ID.Bytes(),
protobuf.MessageType_META,
inputData.NameSpace,
inputData.Project,
inputData.JobID,
inputData.FileName,
inputData.FileHash)
if err != nil {
panic(err)
}

// Send Metadata message
log.Info().Msg("Sending MetaData Packet")
messaging.SendProtoBufMessage(connection, message)

return fileInfo, err
}

//func parseAndSendPayload(bytes *[]byte, numBytes int) {
func parseAndSendPayload(connection net.Conn, inputData *InputData) {

f, err := os.Open(inputData.FilePath)
defer f.Close()
if err != nil {
log.Error().Msgf("Cannot open file %s", inputData.FilePath)
}

log.Info().Msg("Started sending Payload Packets...")
reader := bufio.NewReader(f)
isPayloadFinal := false

buffer := make([]byte, 1024)
for {
n, ioErr := reader.Read(buffer)
if ioErr == io.EOF {
isPayloadFinal = true
// Send terminating payload protobuf message
terminalMessage, err := messaging.CreateMessage_FilePackage(
inputData.ID.Bytes(),
protobuf.MessageType_PACKAGE,
make([]byte, 1),
isPayloadFinal,
)
if err != nil {
log.Fatal().Msg("Fatal Error during payload protobuf assembly")
}
messaging.SendProtoBufMessage(connection, terminalMessage)
break
}

message, err := messaging.CreateMessage_FilePackage(
inputData.ID.Bytes(),
protobuf.MessageType_PACKAGE,
(buffer)[:n],
isPayloadFinal)

if err != nil {
log.Fatal().Msg("Fatal Error during payload protobuf assembly")
}

messaging.SendProtoBufMessage(connection, message)
}
log.Info().Msg("Finished sending Payload Packets...")
}

// Given a valid file path, returns a SHA256 hash
func hashFile(path string) (hash []byte) {
f, err := os.Open(path)
defer f.Close()
if err != nil {
log.Err(err).Msgf("Cannot open file %s", path)
}

hasher := sha256.New()
if _, err := io.Copy(hasher, f); err != nil {
log.Err(err).Msg("Error calculating SHA256 Hash")
}
return hasher.Sum(nil)
}

func listenForACK(connection net.Conn, inputData *InputData) {

// TODO: Make into const in message_handler (also server.go)
// sizePrefix is 4 bytes protobug message size
sizePrefix := make([]byte, 4)

_, _ = io.ReadFull(connection, sizePrefix)
protoLength := int(binary.BigEndian.Uint32(sizePrefix))

ackPacketRaw := make([]byte, protoLength)
_, _ = io.ReadFull(connection, ackPacketRaw)
genericProto := &protobuf.LockerMessage{}
if err := proto.Unmarshal(ackPacketRaw, genericProto); err != nil {
log.Err(err).Msg("Error during unmarshalling")
}

if genericProto.GetAck().ProtoReflect().IsValid() {
ackPacket := genericProto.GetAck()

serverResult := ackPacket.GetServerSuccess()
ackID, _ := xid.FromBytes(ackPacket.GetId())
if ackID != inputData.ID {
log.Warn().
Str("respone_id", ackID.String()).
Str("original_id", inputData.ID.String()).
Msg("Response ID mismatch.")
}

log.Info().
Str("id_back", ackID.String()).
Msgf("ACK packet recieved from server with success flag: %v", serverResult)
}
}

// Parse raw CLI input parameters to internal data structures
func parseInputArguments() []*InputData {
var err error
var inputPath string
const inputPathSeparator = ","

if len(InputArgPath) == 0 {
err = errors.New("--file empty")
log.Err(err).Str("agent", "parseInputArguments").Msgf("No input file was given.")
}
if len(InputArgNamespace) == 0 {
err = errors.New("--namespace empty")
log.Err(err).Str("agent", "parseInputArguments").Msgf("No input namespace was given.")
}
if err != nil {
os.Exit(1)
}

//
// Store information about Namespace, Project and Job-ID
fullNameSpace := InputArgNamespace
namePaths := strings.Split(fullNameSpace, "/")
if len(namePaths) != 3 {
err = errors.New("Namespace must contain 3 values separated by '/'")
log.Err(err).Msg("Namespace values not valid")
os.Exit(1)
}

//
// dataArray contains an *InputData, len(inputPathSlice) times
inputPathSlice := strings.Split(InputArgPath, inputPathSeparator)
dataArray := make([]*InputData, len(inputPathSlice))

for i, path := range inputPathSlice {
if !filepath.IsAbs(path) {
cwd, err := os.Getwd()
if err != nil {
log.Err(err).Msg("Error during CWD PATH parsing")
os.Exit(1)
}
log.Debug().Msgf("Relative path of input: %s", path)
inputPath = filepath.Join(cwd, path)

}

dataArray[i] = &InputData{
FilePath: inputPath,
NameSpace: namePaths[0],
Project: namePaths[1],
JobID: namePaths[2],
ID: xid.New(),
}
}

return dataArray
}

0 comments on commit c87ae22

Please sign in to comment.