From e1d89e86c768297089d4c0a84b0adc2e98493504 Mon Sep 17 00:00:00 2001
From: jatin <84621253+h0x0er@users.noreply.github.com>
Date: Thu, 6 Jun 2024 09:12:19 +0530
Subject: [PATCH] Add one-time-key logic (#430)

---
 agent.go     | 3 ++-
 apiclient.go | 2 ++
 config.go    | 3 +++
 3 files changed, 7 insertions(+), 1 deletion(-)

diff --git a/agent.go b/agent.go
index 80b5cd1..e182fa4 100644
--- a/agent.go
+++ b/agent.go
@@ -65,8 +65,9 @@ func Run(ctx context.Context, configFilePath string, hostDNSServer DNSServer,
 		return err
 	}
 
-	apiclient := &ApiClient{Client: &http.Client{Timeout: 3 * time.Second}, APIURL: config.APIURL, DisableTelemetry: config.DisableTelemetry, EgressPolicy: config.EgressPolicy}
+	apiclient := &ApiClient{Client: &http.Client{Timeout: 3 * time.Second}, APIURL: config.APIURL, DisableTelemetry: config.DisableTelemetry, EgressPolicy: config.EgressPolicy, OneTimeKey: config.OneTimeKey}
 
+	config.OneTimeKey = ""
 	// TODO: pass in an iowriter/ use log library
 	WriteLog(fmt.Sprintf("read config \n %+v", config))
 	WriteLog("\n")
diff --git a/apiclient.go b/apiclient.go
index 2ed19eb..471fdd1 100644
--- a/apiclient.go
+++ b/apiclient.go
@@ -40,6 +40,7 @@ type ApiClient struct {
 	APIURL           string
 	DisableTelemetry bool
 	EgressPolicy     string
+	OneTimeKey       string
 }
 
 const agentApiBaseUrl = "https://apiurl/v1"
@@ -113,6 +114,7 @@ func (apiclient *ApiClient) sendApiRequest(method, url string, body interface{})
 		return err
 	}
 
+	req.Header.Add("x-one-time-key", apiclient.OneTimeKey)
 	if body != nil {
 		req.Header.Add("Content-Type", "application/json; charset=UTF-8")
 	}
diff --git a/config.go b/config.go
index 47f2c90..f927477 100644
--- a/config.go
+++ b/config.go
@@ -16,6 +16,7 @@ type config struct {
 	RunId                 string
 	WorkingDirectory      string
 	APIURL                string
+	OneTimeKey            string
 	Endpoints             map[string][]Endpoint
 	EgressPolicy          string
 	DisableTelemetry      bool
@@ -35,6 +36,7 @@ type configFile struct {
 	RunId                 string `json:"run_id"`
 	WorkingDirectory      string `json:"working_directory"`
 	APIURL                string `json:"api_url"`
+	OneTimeKey            string `json:"one_time_key"`
 	AllowedEndpoints      string `json:"allowed_endpoints"`
 	EgressPolicy          string `json:"egress_policy"`
 	DisableTelemetry      bool   `json:"disable_telemetry"`
@@ -67,6 +69,7 @@ func (c *config) init(configFilePath string) error {
 	c.DisableSudo = configFile.DisableSudo
 	c.DisableFileMonitoring = configFile.DisableFileMonitoring
 	c.Private = configFile.Private
+	c.OneTimeKey = configFile.OneTimeKey
 	return nil
 }