add a function to comment field for publickey struct to detect commen…

…ts, fix test for comments Signed-off-by: Xeckt <jordansavell1198@gmail.com>
opiproject · Jun 21, 2024 · 4e7a4f0 · 4e7a4f0
1 parent 0177a96
commit 4e7a4f0
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 7 deletions.
diff --git a/sztp-agent/pkg/secureagent/utils.go b/sztp-agent/pkg/secureagent/utils.go
@@ -185,10 +185,18 @@ func readSSHHostKeyPublicFiles(pattern string) []publicKey {
 				"Check the key file has the correct format", f, err.Error())
 			continue
 		}
+
+		keyParts := strings.Fields(string(data))
+
 		results = append(results, publicKey{
-			Type:    key.Type(),
-			Data:    strings.Fields(string(data))[1],
-			Comment: "",
+			Type: key.Type(),
+			Data: keyParts[1],
+			Comment: func() string {
+				if len(keyParts) == 3 {
+					return keyParts[2]
+				}
+				return ""
+			}(),
 		})
 	}
 	return results

diff --git a/sztp-agent/pkg/secureagent/utils_test.go b/sztp-agent/pkg/secureagent/utils_test.go
@@ -141,7 +141,8 @@ func Test_readSSHHostKeyPublicFiles(t *testing.T) {
 				content: "ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAID0mjQXlOvkM2HO5vTrSOdHOl3BGOqDiHrx8yYdbP8xR",
 				keyType: "ssh-ed25519",
 			},
-			want: []publicKey{{Type: "ssh-ed25519", Data: "AAAAC3NzaC1lZDI1NTE5AAAAID0mjQXlOvkM2HO5vTrSOdHOl3BGOqDiHrx8yYdbP8xR"}},
+			want: []publicKey{{Type: "ssh-ed25519",
+				Data: "AAAAC3NzaC1lZDI1NTE5AAAAID0mjQXlOvkM2HO5vTrSOdHOl3BGOqDiHrx8yYdbP8xR"}},
 		},
 		{
 			name: "Test OK line in files with comment",
@@ -150,7 +151,9 @@ func Test_readSSHHostKeyPublicFiles(t *testing.T) {
 				content: "ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAID0mjQXlOvkM2HO5vTrSOdHOl3BGOqDiHrx8yYdbP8xR comment",
 				keyType: "ssh-ed25519",
 			},
-			want: []publicKey{{Type: "ssh-ed25519", Data: "AAAAC3NzaC1lZDI1NTE5AAAAID0mjQXlOvkM2HO5vTrSOdHOl3BGOqDiHrx8yYdbP8xR"}},
+			want: []publicKey{{Type: "ssh-ed25519",
+				Data:    "AAAAC3NzaC1lZDI1NTE5AAAAID0mjQXlOvkM2HO5vTrSOdHOl3BGOqDiHrx8yYdbP8xR",
+				Comment: "comment"}},
 		},
 		{
 			name: "Test too many parts in file",
@@ -159,7 +162,8 @@ func Test_readSSHHostKeyPublicFiles(t *testing.T) {
 				content: "ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAID0mjQXlOvkM2HO5vTrSOdHOl3BGOqDiHrx8yYdbP8xR comment error",
 				keyType: "ssh-ed25519",
 			},
-			want: []publicKey{{Type: "ssh-ed25519", Data: "AAAAC3NzaC1lZDI1NTE5AAAAID0mjQXlOvkM2HO5vTrSOdHOl3BGOqDiHrx8yYdbP8xR"}},
+			want: []publicKey{{Type: "ssh-ed25519",
+				Data: "AAAAC3NzaC1lZDI1NTE5AAAAID0mjQXlOvkM2HO5vTrSOdHOl3BGOqDiHrx8yYdbP8xR"}},
 		},
 		{
 			name: "Test not enough parts in file",
@@ -184,7 +188,7 @@ func Test_readSSHHostKeyPublicFiles(t *testing.T) {
 				createTempTestFile(tt.args.file, tt.args.content, true)
 			}
 			if got := readSSHHostKeyPublicFiles(tt.args.file); !reflect.DeepEqual(got, tt.want) {
-				t.Errorf("readSSHHostKeyPublicFiles() = %v, want %v", got, tt.want)
+				t.Errorf("readSSHHostKeyPublicFiles() - got: %v, want %v", got, tt.want)
 			}
 			deleteTempTestFile(tt.args.file)
 		})