124 lines
3.1 KiB
Go
124 lines
3.1 KiB
Go
// NOTE: Most of the code snippet was generated using ChatGPT
|
||
// Prompt used: "generate go program to read and populate ssh authorization file"
|
||
|
||
package p2p
|
||
|
||
import (
|
||
"bufio"
|
||
"errors"
|
||
"fmt"
|
||
"os"
|
||
"path/filepath"
|
||
"strings"
|
||
)
|
||
|
||
// GetAuthorizedKeysPath returns the path to the authorized_keys file
|
||
func GetAuthorizedKeysPath() (string, error) {
|
||
homeDir, err := os.UserHomeDir()
|
||
if err != nil {
|
||
return "", fmt.Errorf("could not find home directory: %v", err)
|
||
}
|
||
return filepath.Join(homeDir, ".ssh", "authorized_keys"), nil
|
||
}
|
||
|
||
// ReadAuthorizedKeys reads and returns the current contents of the authorized_keys file as a map
|
||
func ReadAuthorizedKeys(path string) (map[string]bool, error) {
|
||
file, err := os.Open(path)
|
||
if err != nil {
|
||
return nil, fmt.Errorf("could not open authorized_keys file: %v", err)
|
||
}
|
||
defer file.Close()
|
||
|
||
keys := make(map[string]bool)
|
||
scanner := bufio.NewScanner(file)
|
||
for scanner.Scan() {
|
||
line := strings.TrimSpace(scanner.Text())
|
||
// Skip empty lines and comments
|
||
if line != "" && !strings.HasPrefix(line, "#") {
|
||
keys[line] = true
|
||
}
|
||
}
|
||
if err := scanner.Err(); err != nil {
|
||
return nil, fmt.Errorf("error reading authorized_keys file: %v", err)
|
||
}
|
||
return keys, nil
|
||
}
|
||
|
||
// AddKeyToAuthorizedKeys adds a new key to the authorized_keys file if it doesn’t already exist
|
||
func AddKeyToAuthorizedKeys(path, newKey string) error {
|
||
keys, err := ReadAuthorizedKeys(path)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
|
||
// Check if the key already exists in the map
|
||
if keys[newKey] {
|
||
return errors.New("key already exists in authorized_keys")
|
||
}
|
||
|
||
// Append the new key
|
||
file, err := os.OpenFile(path, os.O_APPEND|os.O_WRONLY|os.O_CREATE, 0600)
|
||
if err != nil {
|
||
return fmt.Errorf("could not open authorized_keys file for writing: %v", err)
|
||
}
|
||
defer file.Close()
|
||
|
||
if _, err := file.WriteString(newKey + "\n"); err != nil {
|
||
return fmt.Errorf("could not write to authorized_keys file: %v", err)
|
||
}
|
||
return nil
|
||
}
|
||
|
||
func RemoveKeyFromAuthorizedKeys(path, keyToRemove string) error {
|
||
keys, err := ReadAuthorizedKeys(path)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
|
||
// Check if the key exists in the map
|
||
if !keys[keyToRemove] {
|
||
return errors.New("key not found in authorized_keys")
|
||
}
|
||
|
||
// Delete the key from the map
|
||
delete(keys, keyToRemove)
|
||
|
||
// Write updated keys back to the authorized_keys file
|
||
file, err := os.OpenFile(path, os.O_TRUNC|os.O_WRONLY|os.O_CREATE, 0600)
|
||
if err != nil {
|
||
return fmt.Errorf("could not open authorized_keys file for writing: %v", err)
|
||
}
|
||
defer file.Close()
|
||
|
||
for key := range keys {
|
||
if _, err := file.WriteString(key + "\n"); err != nil {
|
||
return fmt.Errorf("could not write to authorized_keys file: %v", err)
|
||
}
|
||
}
|
||
|
||
return nil
|
||
}
|
||
|
||
// AddAuthorisationKey Adds public key provided to the
|
||
// authorization file so that nodes can SSH into
|
||
// the
|
||
func AddAuthorisationKey(PublicKey string) error {
|
||
path, err := GetAuthorizedKeysPath()
|
||
if err != nil {
|
||
return err
|
||
}
|
||
|
||
// Display existing keys
|
||
_, err = ReadAuthorizedKeys(path)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
|
||
err = AddKeyToAuthorizedKeys(path, PublicKey)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
|
||
return nil
|
||
}
|