iter23

cmd/agent/main.go

@@ -29,6 +29,10 @@ func main() {
  log.Fatal(err)
  }
 
- cfg := config.NewAgentConfig()
+ cfg, err := config.NewAgentConfig()
+ if err != nil {
+ log.Fatal(err)
+ }
+
  agent.StartClient(ctx, cfg)
 }

cmd/key/main.go

@@ -0,0 +1,45 @@
+package main
+
+import (
+ "crypto/rand"
+ "crypto/rsa"
+ "crypto/x509"
+ "encoding/pem"
+ "log"
+ "os"
+)
+
+func main() {
+ privateKey, err := rsa.GenerateKey(rand.Reader, 2048)
+ if err != nil {
+ log.Fatal(err)
+ }
+
+ publicKey := &privateKey.PublicKey
+ privateKeyBytes := x509.MarshalPKCS1PrivateKey(privateKey)
+
+ privateKeyPEM := pem.EncodeToMemory(&pem.Block{
+ Type: "SERVER PRIVATE KEY",
+ Bytes: privateKeyBytes,
+ })
+
+ err = os.WriteFile("./private.pem", privateKeyPEM, 0644)
+ if err != nil {
+ log.Fatal(err)
+ }
+
+ publicKeyBytes, err := x509.MarshalPKIXPublicKey(publicKey)
+ if err != nil {
+ log.Fatal(err)
+ }
+
+ publicKeyPEM := pem.EncodeToMemory(&pem.Block{
+ Type: "AGENT PUBLIC KEY",
+ Bytes: publicKeyBytes,
+ })
+
+ err = os.WriteFile("./public.pem", publicKeyPEM, 0644)
+ if err != nil {
+ log.Fatal(err)
+ }
+}

cmd/server/main.go

@@ -19,6 +19,9 @@ func main() {
  log.Fatal(err)
  }
 
- cfg := config.NewServerConfig()
+ cfg, err := config.NewServerConfig()
+ if err != nil {
+ log.Fatal(err)
+ }
  server.StartListener(cfg)
 }

internal/pkg/agent/config/config.go

@@ -1,18 +1,26 @@
 package config
 
 import (
+ "crypto/rsa"
+ "crypto/x509"
+ "encoding/json"
+ "encoding/pem"
  "flag"
+ "os"
 
  "github.com/caarlos0/env/v6"
  "github.com/sirupsen/logrus"
 )
 
 type AgentConfig struct {
- ServerAddress string `env:"ADDRESS"`
+ ServerAddress string `env:"ADDRESS" json:"server_address"`
  SignKey string `env:"KEY"`
- ReportInterval int `env:"REPORT_INTERVAL"`
- PollInterval int `env:"POLL_INTERVAL"`
+ ReportInterval int `env:"REPORT_INTERVAL" json:"report_interval"`
+ PollInterval int `env:"POLL_INTERVAL" json:"poll_interval"`
  RateLimit int `env:"RATE_LIMIT"`
+ PublicKeyPath string `env:"CRYPTO_KEY" json:"crypto_key"`
+ PublicKey *rsa.PublicKey
+ ConfigPath string `env:"CONFIG"`
  SignKeyByte []byte
 }
 
@@ -23,7 +31,7 @@ const (
  rateLimitDefault = 3
 )
 
-func NewAgentConfig() *AgentConfig {
+func NewAgentConfig() (*AgentConfig, error) {
  cfg := AgentConfig{}
  cfg.init()
 
@@ -33,9 +41,25 @@ func NewAgentConfig() *AgentConfig {
 
  if err := env.Parse(&cfg); err != nil {
  logrus.Errorf("env parsing error: %v", err)
- return nil
+ return nil, err
  }
- return &cfg
+
+ if cfg.ConfigPath != "" {
+ cfgJSON, err := readConfigFile(cfg.ConfigPath)
+ if err != nil {
+ return cfgJSON, err
+ }
+ }
+
+ if cfg.PublicKeyPath != "" {
+ publicKey, err := cfg.getPublicKey()
+ if err != nil {
+ logrus.Errorf("error with get public key: %v", err)
+ }
+ cfg.PublicKey = publicKey
+ }
+
+ return &cfg, nil
 }
 
 func (c *AgentConfig) init() {
@@ -44,5 +68,41 @@ func (c *AgentConfig) init() {
  flag.IntVar(&c.PollInterval, "p", pollIntervalDefault, "Interval of poll metric")
  flag.StringVar(&c.SignKey, "k", "", "Server key")
  flag.IntVar(&c.RateLimit, "l", rateLimitDefault, "Rate limit")
+ flag.StringVar(&c.PublicKeyPath, "-crypto-key", "", "Public key path")
+ flag.StringVar(&c.ConfigPath, "c", "", "Path to config file")
+ flag.StringVar(&c.ConfigPath, "config", "", "Path to config file (the same as -c)")
  flag.Parse()
 }
+
+func readConfigFile(path string) (cfg *AgentConfig, err error) {
+ data, err := os.ReadFile(path)
+ if err != nil {
+ return cfg, err
+ }
+
+ err = json.Unmarshal(data, &cfg)
+ if err != nil {
+ return cfg, err
+ }
+
+ return cfg, nil
+}
+
+func (c *AgentConfig) getPublicKey() (*rsa.PublicKey, error) {
+ publicKeyPEM, err := os.ReadFile(c.PublicKeyPath)
+ if err != nil {
+ return nil, err
+ }
+
+ publicKeyBlock, _ := pem.Decode(publicKeyPEM)
+ if publicKeyBlock == nil {
+ return nil, err
+ }
+
+ publicKey, err := x509.ParsePKCS1PublicKey(publicKeyBlock.Bytes)
+ if err != nil {
+ return nil, err
+ }
+
+ return publicKey, nil
+}

internal/pkg/agent/config/config_test.go

@@ -39,7 +39,7 @@ func TestNewAgentConfig(t *testing.T) {
  }
  for _, tt := range tests {
  t.Run(tt.name, func(t *testing.T) {
- if got := NewAgentConfig(); !reflect.DeepEqual(got, tt.want) {
+ if got, _ := NewAgentConfig(); !reflect.DeepEqual(got, tt.want) {
  t.Errorf("NewAgentConfig() = %v, want %v", got, tt.want)
  }
  })

internal/pkg/agent/encrypt.go

@@ -0,0 +1,15 @@
+package agent
+
+import (
+ "crypto/rand"
+ "crypto/rsa"
+)
+
+func encrypt(key *rsa.PublicKey, data []byte) ([]byte, error) {
+ ciphertext, err := rsa.EncryptPKCS1v15(rand.Reader, key, data)
+ if err != nil {
+ return nil, err
+ }
+
+ return ciphertext, nil
+}

internal/pkg/agent/reporter.go

@@ -24,7 +24,7 @@ func RunSendMetric(ctx context.Context, reportTicker *time.Ticker, c *config.Age
  case <-ctx.Done():
  return
  case <-reportTicker.C:
- ok, err := SendMetrics(ctx, s, c.ServerAddress, c.SignKeyByte)
+ ok, err := SendMetrics(ctx, s, c)
  if err != nil {
  logrus.Errorf("Error send metrics %v", err)
  }
@@ -37,7 +37,7 @@ func RunSendMetric(ctx context.Context, reportTicker *time.Ticker, c *config.Age
  }
 }
 
-func SendMetrics(ctx context.Context, s storage.Store, serverAddress string, signKey []byte) (bool, error) {
+func SendMetrics(ctx context.Context, s storage.Store, c *config.AgentConfig) (bool, error) {
  metricsMap, err := s.GetMetrics(ctx)
  if err != nil {
  logrus.Errorf("Some error ocured during metrics get: %q", err)
@@ -49,20 +49,27 @@ func SendMetrics(ctx context.Context, s storage.Store, serverAddress string, sig
  metricsBatch = append(metricsBatch, v)
  }
 
- url := fmt.Sprintf("http://%s/updates/", serverAddress)
+ url := fmt.Sprintf("http://%s/updates/", c.ServerAddress)
 
- if err = SendBatchJSON(url, metricsBatch, signKey); err != nil {
+ if err = SendBatchJSON(url, metricsBatch, c); err != nil {
  return false, fmt.Errorf("error create post request %w", err)
  }
  return true, nil
 }
 
-func SendBatchJSON(url string, metricsBatch []*metrics.Metrics, signKey []byte) error {
+func SendBatchJSON(url string, metricsBatch []*metrics.Metrics, c *config.AgentConfig) error {
  body, err := json.Marshal(metricsBatch)
  if err != nil {
  return fmt.Errorf("error encoding metric %w", err)
  }
 
+ if c.PublicKeyPath != "" {
+ body, err = encrypt(c.PublicKey, body)
+ if err != nil {
+ return err
+ }
+ }
+
  var buf bytes.Buffer
  gz := gzip.NewWriter(&buf)
  if _, err = gz.Write(body); err != nil {
@@ -82,8 +89,8 @@ func SendBatchJSON(url string, metricsBatch []*metrics.Metrics, signKey []byte)
  req.Header.Set("Content-Type", "application/json")
  req.Header.Set("Content-Encoding", "gzip")
 
- if signKey != nil {
- h := hmac.New(sha256.New, signKey)
+ if c.SignKeyByte != nil {
+ h := hmac.New(sha256.New, c.SignKeyByte)
  h.Write(body)
  serverHash := hex.EncodeToString(h.Sum(nil))
  req.Header.Set("HashSHA256", serverHash)

internal/pkg/agent/reporter_test.go

@@ -10,6 +10,7 @@ import (
  "testing"
 
  "github.com/mayr0y/animated-octo-couscous.git/internal/pkg/agent"
+ "github.com/mayr0y/animated-octo-couscous.git/internal/pkg/agent/config"
  "github.com/mayr0y/animated-octo-couscous.git/internal/pkg/metrics"
  "github.com/mayr0y/animated-octo-couscous.git/internal/pkg/storage"
 )
@@ -22,7 +23,7 @@ const (
 )
 
 func TestSendReport(t *testing.T) {
- secretKey := []byte("secret")
+ c, _ := config.NewAgentConfig()
  mtr := storage.NewMetrics()
 
  err := agent.UpdateMetrics(context.Background(), mtr)
@@ -70,5 +71,5 @@ func TestSendReport(t *testing.T) {
  }))
  defer server.Close()
 
- agent.SendMetrics(context.Background(), mtr, server.URL, secretKey)
+ agent.SendMetrics(context.Background(), mtr, c)
 }

internal/pkg/middleware/crypt.go

@@ -1,7 +1,10 @@
 package middleware
 
 import (
+ "bytes"
  "crypto/hmac"
+ "crypto/rand"
+ "crypto/rsa"
  "crypto/sha256"
  "encoding/hex"
  "io"
@@ -51,3 +54,24 @@ func CryptMiddleware(signKey []byte) func(handler http.Handler) http.Handler {
  })
  }
 }
+
+func DecryptMiddleware(key *rsa.PrivateKey) func(http.Handler) http.Handler {
+ return func(next http.Handler) http.Handler {
+ return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+ body, err := io.ReadAll(r.Body)
+ if err != nil {
+ http.Error(w, err.Error(), http.StatusBadRequest)
+ return
+ }
+
+ plaintext, err := rsa.DecryptPKCS1v15(rand.Reader, key, body)
+ if err != nil {
+ http.Error(w, err.Error(), http.StatusBadRequest)
+ return
+ }
+
+ r.Body = io.NopCloser(bytes.NewBuffer(plaintext))
+ next.ServeHTTP(w, r)
+ })
+ }
+}