p2p-llm / config /config.go
arpinfidel's picture
temp
48511d8
package config
import (
"crypto/rand"
"crypto/rsa"
"crypto/x509"
"encoding/json"
"encoding/pem"
"fmt"
"os"
"strings"
)
type Peer struct {
URL string
PublicKey *rsa.PublicKey
}
type Config struct {
DBPath string `json:"db_path"`
TargetURL string `json:"target_url"`
Port string `json:"port"`
MaxParallelRequests int `json:"max_parallel_requests"`
TrustedPeers []Peer `json:"trusted_peers"`
TrustedPeersPath string `json:"trusted_peers_path"`
}
type Secrets struct {
PrivateKeyPath string `json:"jwt_private_key_path"`
PublicKeyPath string `json:"jwt_public_key_path"`
}
type KeyPair struct {
PrivateKey *rsa.PrivateKey
PublicKey *rsa.PublicKey
}
func LoadConfig(path string) (*Config, error) {
file, err := os.ReadFile(path)
if err != nil {
return nil, err
}
var cfg Config
err = json.Unmarshal(file, &cfg)
return &cfg, err
}
func LoadSecrets(path string) (*Secrets, error) {
file, err := os.ReadFile(path)
if err != nil {
return nil, err
}
var secrets Secrets
err = json.Unmarshal(file, &secrets)
return &secrets, err
}
func generateRSAKeys() (*rsa.PrivateKey, error) {
return rsa.GenerateKey(rand.Reader, 2048)
}
func saveKeyToFile(key interface{}, path string) error {
file, err := os.Create(path)
if err != nil {
return err
}
defer file.Close()
var pemBlock *pem.Block
switch k := key.(type) {
case *rsa.PrivateKey:
bytes, err := x509.MarshalPKCS8PrivateKey(k)
if err != nil {
return err
}
pemBlock = &pem.Block{
Type: "PRIVATE KEY",
Bytes: bytes,
}
case *rsa.PublicKey:
bytes, err := x509.MarshalPKIXPublicKey(k)
if err != nil {
return err
}
pemBlock = &pem.Block{
Type: "PUBLIC KEY",
Bytes: bytes,
}
default:
return fmt.Errorf("unsupported key type")
}
return pem.Encode(file, pemBlock)
}
func loadKeyFromFile(path string, private bool) (interface{}, error) {
data, err := os.ReadFile(path)
if err != nil {
return nil, err
}
block, _ := pem.Decode(data)
if block == nil {
return nil, fmt.Errorf("failed to parse PEM block")
}
if private {
key, err := x509.ParsePKCS8PrivateKey(block.Bytes)
if err != nil {
return nil, err
}
return key.(*rsa.PrivateKey), nil
}
return x509.ParsePKIXPublicKey(block.Bytes)
}
func ensureKeysExist(secrets *Secrets) (*KeyPair, error) {
if _, err := os.Stat(secrets.PrivateKeyPath); os.IsNotExist(err) {
privateKey, err := generateRSAKeys()
if err != nil {
return nil, err
}
if err := saveKeyToFile(privateKey, secrets.PrivateKeyPath); err != nil {
return nil, err
}
if err := saveKeyToFile(&privateKey.PublicKey, secrets.PublicKeyPath); err != nil {
return nil, err
}
}
privateKey, err := loadKeyFromFile(secrets.PrivateKeyPath, true)
if err != nil {
return nil, err
}
publicKey, err := loadKeyFromFile(secrets.PublicKeyPath, false)
if err != nil {
return nil, err
}
return &KeyPair{
PrivateKey: privateKey.(*rsa.PrivateKey),
PublicKey: publicKey.(*rsa.PublicKey),
}, nil
}
func LoadTrustedPeers(path string) ([]Peer, error) {
data, err := os.ReadFile(path)
if err != nil {
return nil, err
}
var peers []Peer
lines := strings.Split(string(data), "\n")
for _, line := range lines {
line = strings.TrimSpace(line)
if line == "" || strings.HasPrefix(line, "#") {
continue
}
parts := strings.SplitN(line, "|", 2)
if len(parts) != 2 {
continue
}
block, _ := pem.Decode([]byte(parts[1]))
if block == nil {
continue
}
pubKey, err := x509.ParsePKIXPublicKey(block.Bytes)
if err != nil {
continue
}
peers = append(peers, Peer{
URL: parts[0],
PublicKey: pubKey.(*rsa.PublicKey),
})
}
return peers, nil
}
func Load() (*Config, *Secrets, *KeyPair, error) {
cfg, err := LoadConfig("/app/files/config.json")
if err != nil {
return nil, nil, nil, err
}
secrets, err := LoadSecrets("/app/files/secrets.json")
if err != nil {
return nil, nil, nil, err
}
keyPair, err := ensureKeysExist(secrets)
if err != nil {
return nil, nil, nil, err
}
if cfg.TrustedPeersPath != "" {
peers, err := LoadTrustedPeers(cfg.TrustedPeersPath)
if err != nil {
return nil, nil, nil, err
}
cfg.TrustedPeers = peers
}
return cfg, secrets, keyPair, nil
}