From 911e080fbbdf570e88244bdd0b85f7c1fbb57fcc Mon Sep 17 00:00:00 2001
From: Myzel394 <50424412+Myzel394@users.noreply.github.com>
Date: Sun, 18 Aug 2024 22:14:19 +0200
Subject: [PATCH] feat(wireguard): Add code actions

---
 handlers/wireguard/code-actions.go            | 51 +++++++++++++++++--
 .../wireguard/text-document-code-action.go    | 23 +++++++++
 handlers/wireguard/wg-commands.go             | 19 ++++++-
 handlers/wireguard/wg-parser.go               |  2 +-
 .../wireguard/workspace-execute-command.go    |  6 +++
 5 files changed, 94 insertions(+), 7 deletions(-)

diff --git a/handlers/wireguard/code-actions.go b/handlers/wireguard/code-actions.go
index 1894130..9afdc19 100644
--- a/handlers/wireguard/code-actions.go
+++ b/handlers/wireguard/code-actions.go
@@ -7,7 +7,8 @@ import (
 type codeActionName string
 
 const (
-	codeActionGeneratePrivateKey codeActionName = "generatePrivateKey"
+	codeActionGeneratePrivateKey   codeActionName = "generatePrivateKey"
+	codeActionGeneratePresharedKey codeActionName = "generatePresharedKey"
 )
 
 type codeActionGeneratePrivateKeyArgs struct {
@@ -22,13 +23,26 @@ func codeActionGeneratePrivateKeyArgsFromArguments(arguments map[string]any) cod
 	}
 }
 
+type codeActionGeneratePresharedKeyArgs struct {
+	URI  protocol.DocumentUri
+	Line uint32
+}
+
+func codeActionGeneratePresharedKeyArgsFromArguments(arguments map[string]any) codeActionGeneratePresharedKeyArgs {
+	return codeActionGeneratePresharedKeyArgs{
+		URI:  arguments["URI"].(protocol.DocumentUri),
+		Line: uint32(arguments["Line"].(float64)),
+	}
+}
+
 func (p wireguardProperty) getInsertRange(line uint32) protocol.Range {
-	insertPosition := p.Separator.Location.End
+	var insertPosition uint32 = p.Separator.Location.End
 	var length uint32 = 0
 
 	if p.Value != nil {
-		// Length of the value
-		length = p.Value.Location.End - p.Value.Location.Start
+		insertPosition = p.Value.Location.Start - 1
+		// Length of the value; +1 because of the starting space
+		length = (p.Value.Location.End - p.Value.Location.Start) + 1
 	}
 
 	return protocol.Range{
@@ -71,3 +85,32 @@ func (p *wireguardParser) runGeneratePrivateKey(args codeActionGeneratePrivateKe
 		},
 	}, nil
 }
+
+func (p *wireguardParser) runGeneratePresharedKey(args codeActionGeneratePresharedKeyArgs) (*protocol.ApplyWorkspaceEditParams, error) {
+	presharedKey, err := createPresharedKey()
+
+	if err != nil {
+		return &protocol.ApplyWorkspaceEditParams{}, err
+	}
+
+	section, property := p.getPropertyByLine(args.Line)
+
+	if section == nil || property == nil {
+		return nil, nil
+	}
+
+	label := "Generate Preshared Key"
+	return &protocol.ApplyWorkspaceEditParams{
+		Label: &label,
+		Edit: protocol.WorkspaceEdit{
+			Changes: map[protocol.DocumentUri][]protocol.TextEdit{
+				args.URI: {
+					{
+						NewText: " " + presharedKey,
+						Range:   property.getInsertRange(args.Line),
+					},
+				},
+			},
+		},
+	}, nil
+}
diff --git a/handlers/wireguard/text-document-code-action.go b/handlers/wireguard/text-document-code-action.go
index 892f2bf..40b9a71 100644
--- a/handlers/wireguard/text-document-code-action.go
+++ b/handlers/wireguard/text-document-code-action.go
@@ -39,6 +39,29 @@ func TextDocumentCodeAction(context *glsp.Context, params *protocol.CodeActionPa
 				Command: &command,
 			},
 		}, nil
+	case "PresharedKey":
+		if !areWireguardToolsAvailable() {
+			return nil, nil
+		}
+
+		commandID := "wireguard." + codeActionGeneratePresharedKey
+		command := protocol.Command{
+			Title:   "Generate PresharedKey",
+			Command: string(commandID),
+			Arguments: []any{
+				codeActionGeneratePresharedKeyArgs{
+					URI:  params.TextDocument.URI,
+					Line: line,
+				},
+			},
+		}
+
+		return []protocol.CodeAction{
+			{
+				Title:   "Generate PresharedKey",
+				Command: &command,
+			},
+		}, nil
 	}
 
 	return nil, nil
diff --git a/handlers/wireguard/wg-commands.go b/handlers/wireguard/wg-commands.go
index 7a345cf..37660dc 100644
--- a/handlers/wireguard/wg-commands.go
+++ b/handlers/wireguard/wg-commands.go
@@ -2,9 +2,12 @@ package wireguard
 
 import (
 	"os/exec"
+	"regexp"
 	"strings"
 )
 
+var whitespacePattern = regexp.MustCompile(`[\s\n]+`)
+
 func areWireguardToolsAvailable() bool {
 	_, err := exec.LookPath("wg")
 
@@ -20,7 +23,7 @@ func createNewPrivateKey() (string, error) {
 		return "", err
 	}
 
-	return string(bytes), nil
+	return string(whitespacePattern.ReplaceAll(bytes, []byte(""))), nil
 }
 
 func createPublicKey(privateKey string) (string, error) {
@@ -33,5 +36,17 @@ func createPublicKey(privateKey string) (string, error) {
 		return "", err
 	}
 
-	return strings.ReplaceAll(string(bytes), "\n", ""), nil
+	return string(whitespacePattern.ReplaceAll(bytes, []byte(""))), nil
+}
+
+func createPresharedKey() (string, error) {
+	cmd := exec.Command("wg", "genpsk")
+
+	bytes, err := cmd.Output()
+
+	if err != nil {
+		return "", err
+	}
+
+	return string(whitespacePattern.ReplaceAll(bytes, []byte(""))), nil
 }
diff --git a/handlers/wireguard/wg-parser.go b/handlers/wireguard/wg-parser.go
index 65f08e0..a035185 100644
--- a/handlers/wireguard/wg-parser.go
+++ b/handlers/wireguard/wg-parser.go
@@ -331,7 +331,7 @@ func (p *wireguardParser) getBelongingSectionByLine(line uint32) *wireguardSecti
 func (p *wireguardParser) getPropertyByLine(line uint32) (*wireguardSection, *wireguardProperty) {
 	section := p.getSectionByLine(line)
 
-	if section.Name == nil {
+	if section == nil || section.Name == nil {
 		return nil, nil
 	}
 
diff --git a/handlers/wireguard/workspace-execute-command.go b/handlers/wireguard/workspace-execute-command.go
index 39a1a60..f803733 100644
--- a/handlers/wireguard/workspace-execute-command.go
+++ b/handlers/wireguard/workspace-execute-command.go
@@ -17,6 +17,12 @@ func WorkspaceExecuteCommand(context *glsp.Context, params *protocol.ExecuteComm
 		parser := documentParserMap[args.URI]
 
 		return parser.runGeneratePrivateKey(args)
+	case string(codeActionGeneratePresharedKey):
+		args := codeActionGeneratePresharedKeyArgsFromArguments(params.Arguments[0].(map[string]any))
+
+		parser := documentParserMap[args.URI]
+
+		return parser.runGeneratePresharedKey(args)
 	}
 
 	return nil, nil