Automatic Python vulnerability insertion

generate_text.go

@@ -24,9 +24,9 @@ func generateCommit(questionText string) string {
 	}
 
 	// printResponse(resp)
-	
-	return string(resp.Candidates[0].Content.Parts[0].(genai.Text)) 
-	
+
+	return string(resp.Candidates[0].Content.Parts[0].(genai.Text))
+
 }
 
 func printResponse(resp *genai.GenerateContentResponse) {

get_vuln.go

@@ -0,0 +1,136 @@
+package main
+
+import (
+	"context"
+	"github.com/google/go-github/v57/github"
+	"log"
+	"math/rand/v2"
+	"os"
+	"path/filepath"
+	"strings"
+)
+
+func Ptr[T any](v T) *T {
+	return &v
+}
+
+func getPythonVuln(githubPersonalAccessToken string) (cve, packageName, version string) {
+	cve, packageName, version = getVuln(githubPersonalAccessToken, "pip")
+	return
+}
+
+func getVuln(githubPersonalAccessToken string, ecosystem string) (cve, packageName, version string) {
+	ctx := context.Background()
+	client := github.NewClient(nil).WithAuthToken(githubPersonalAccessToken)
+
+	options := github.ListGlobalSecurityAdvisoriesOptions{
+		Ecosystem: Ptr(ecosystem),
+		Severity:  Ptr("critical"),
+	}
+
+	securityService := client.SecurityAdvisories
+	vulns, _, err := securityService.ListGlobalSecurityAdvisories(ctx, &options)
+	if err != nil {
+		log.Fatal(err)
+	}
+
+	// Pick one of the first five reports
+	j := rand.IntN(20)
+	// Some reports don't have CVEID, need to skip if that's the case
+	for {
+		log.Println("Evaluating " + vulns[j].GetCVEID() + "  " + vulns[j].Vulnerabilities[0].GetVulnerableVersionRange())
+		if vulns[j].GetCVEID() == "" {
+			j = j + 1
+			continue
+		}
+
+		// Skip versions with < only
+		if r := vulns[j].Vulnerabilities[0].GetVulnerableVersionRange(); strings.Contains(r, "< ") {
+			j = j + 1
+			continue
+		}
+
+		// Want only fixable vulns
+		if r := vulns[j].Vulnerabilities[0].GetFirstPatchedVersion(); r == "" {
+			j = j + 1
+			continue
+		}
+		break
+	}
+	cve = vulns[j].GetCVEID()
+	packageName = vulns[j].Vulnerabilities[0].GetPackage().GetName()
+	version = ""
+	r := vulns[j].Vulnerabilities[0].GetVulnerableVersionRange()
+
+	// Extract version from version range
+	switch {
+	case strings.Contains(r, "<="):
+		_, version, _ = strings.Cut(r, "<= ")
+	/* revisit
+	case strings.Contains(r, "<"):
+		_, version, _ = strings.Cut(r, "< ")
+		s := strings.Split(version, ".")
+		log.Println(s)
+		i, _ := strconv.Atoi(s[2])
+		s[2] = strconv.Itoa(i - 1)
+
+		version = s[0] + "." + s[1] + "." + s[2]
+	*/
+	case strings.Contains(r, "="):
+		_, version, _ = strings.Cut(r, "= ")
+	default:
+		log.Fatal("Unexpected version range string " + r)
+	}
+
+	return
+}
+
+func writePythonVuln(packageName string, version string) {
+	requirementsFile := findRequirementsTxt("./terragoat/")
+	if requirementsFile == "" {
+		return
+	}
+	b, err := os.ReadFile(requirementsFile) // just pass the file name
+	if err != nil {
+		log.Fatal(err)
+	}
+
+	requirementsContents := string(b)
+
+	if strings.HasPrefix(requirementsContents, packageName+" ==") {
+
+		lines := strings.Split(string(requirementsContents), "\n")
+
+		for i, line := range lines {
+			if strings.Contains(line, packageName) {
+				lines[i] = packageName + "==" + version
+			}
+		}
+		requirementsContents = strings.Join(lines, "\n")
+
+	} else {
+		requirementsContents = requirementsContents + "\n" + packageName + "==" + version
+	}
+
+	err = os.WriteFile(requirementsFile, []byte(requirementsContents), 0644)
+	if err != nil {
+		log.Fatal(err)
+	}
+
+}
+
+func findRequirementsTxt(repoRoot string) (requirementsFile string) {
+	err := filepath.Walk(repoRoot, func(path string, info os.FileInfo, err error) error {
+		if err != nil {
+			log.Fatal(err)
+		}
+		if strings.Index(info.Name(), "requirements.txt") != -1 {
+			requirementsFile = path
+		}
+		return nil
+	})
+	if err != nil {
+		log.Fatal(err)
+	}
+	return
+}

go.mod

@@ -39,6 +39,8 @@ require (
 	github.com/kevinburke/ssh_config v1.2.0 // indirect
 	github.com/pjbgf/sha1cd v0.3.0 // indirect
 	github.com/sergi/go-diff v1.3.2-0.20230802210424-5b0b94c5c0d3 // indirect
+	github.com/shurcooL/githubv4 v0.0.0-20240727222349-48295856cce7 // indirect
+	github.com/shurcooL/graphql v0.0.0-20200928012149-18c5c3165e3a // indirect
 	github.com/skeema/knownhosts v1.3.0 // indirect
 	github.com/xanzy/ssh-agent v0.3.3 // indirect
 	go.opencensus.io v0.24.0 // indirect
@@ -47,14 +49,15 @@ require (
 	go.opentelemetry.io/otel v1.26.0 // indirect
 	go.opentelemetry.io/otel/metric v1.26.0 // indirect
 	go.opentelemetry.io/otel/trace v1.26.0 // indirect
-	golang.org/x/crypto v0.31.0 // indirect
-	golang.org/x/mod v0.17.0 // indirect
-	golang.org/x/net v0.33.0 // indirect
+	golang.org/x/crypto v0.32.0 // indirect
+	golang.org/x/mod v0.22.0 // indirect
+	golang.org/x/net v0.34.0 // indirect
 	golang.org/x/sync v0.10.0 // indirect
-	golang.org/x/sys v0.28.0 // indirect
+	golang.org/x/sys v0.29.0 // indirect
 	golang.org/x/text v0.21.0 // indirect
-	golang.org/x/time v0.5.0 // indirect
-	golang.org/x/tools v0.21.1-0.20240508182429-e35e4ccd0d2d // indirect
+	golang.org/x/time v0.9.0 // indirect
+	golang.org/x/tools v0.29.0 // indirect
+	golang.org/x/vulndb v0.0.0-20250130185858-953c816d6fd6 // indirect
 	google.golang.org/api v0.186.0 // indirect
 	google.golang.org/genproto/googleapis/api v0.0.0-20240617180043-68d350f18fd4 // indirect
 	google.golang.org/genproto/googleapis/rpc v0.0.0-20240617180043-68d350f18fd4 // indirect

go.sum

@@ -124,6 +124,10 @@ github.com/rogpeppe/go-internal v1.12.0 h1:exVL4IDcn6na9z1rAb56Vxr+CgyK3nn3O+epU
 github.com/rogpeppe/go-internal v1.12.0/go.mod h1:E+RYuTGaKKdloAfM02xzb0FW3Paa99yedzYV+kq4uf4=
 github.com/sergi/go-diff v1.3.2-0.20230802210424-5b0b94c5c0d3 h1:n661drycOFuPLCN3Uc8sB6B/s6Z4t2xvBgU1htSHuq8=
 github.com/sergi/go-diff v1.3.2-0.20230802210424-5b0b94c5c0d3/go.mod h1:A0bzQcvG0E7Rwjx0REVgAGH58e96+X0MeOfepqsbeW4=
+github.com/shurcooL/githubv4 v0.0.0-20240727222349-48295856cce7 h1:cYCy18SHPKRkvclm+pWm1Lk4YrREb4IOIb/YdFO0p2M=
+github.com/shurcooL/githubv4 v0.0.0-20240727222349-48295856cce7/go.mod h1:zqMwyHmnN/eDOZOdiTohqIUKUrTFX62PNlu7IJdu0q8=
+github.com/shurcooL/graphql v0.0.0-20200928012149-18c5c3165e3a h1:KikTa6HtAK8cS1qjvUvvq4QO21QnwC+EfvB+OAuZ/ZU=
+github.com/shurcooL/graphql v0.0.0-20200928012149-18c5c3165e3a/go.mod h1:AuYgA5Kyo4c7HfUmvRGs/6rGlMMV/6B1bVnB9JxJEEg=
 github.com/sirupsen/logrus v1.7.0/go.mod h1:yWOB1SBYBC5VeMP7gHvWumXLIWorT60ONWic61uBYv0=
 github.com/skeema/knownhosts v1.3.0 h1:AM+y0rI04VksttfwjkSTNQorvGqmwATnvnAHpSgc0LY=
 github.com/skeema/knownhosts v1.3.0/go.mod h1:sPINvnADmT/qYH1kfv+ePMmOBTH6Tbl7b5LvTDjFK7M=
@@ -156,6 +160,8 @@ golang.org/x/crypto v0.0.0-20200622213623-75b288015ac9/go.mod h1:LzIPMQfyMNhhGPh
 golang.org/x/crypto v0.0.0-20220622213112-05595931fe9d/go.mod h1:IxCIyHEi3zRg3s0A5j5BB6A9Jmi73HwBIUl50j+osU4=
 golang.org/x/crypto v0.31.0 h1:ihbySMvVjLAeSH1IbfcRTkD/iNscyz8rGzjF/E5hV6U=
 golang.org/x/crypto v0.31.0/go.mod h1:kDsLvtWBEx7MV9tJOj9bnXsPbxwJQ6csT/x4KIN4Ssk=
+golang.org/x/crypto v0.32.0 h1:euUpcYgM8WcP71gNpTqQCn6rC2t6ULUPiOzfWaXVVfc=
+golang.org/x/crypto v0.32.0/go.mod h1:ZnnJkOaASj8g0AjIduWNlq2NRxL0PlBrbKVyZ6V/Ugc=
 golang.org/x/exp v0.0.0-20190121172915-509febef88a4/go.mod h1:CJ0aWSM057203Lf6IL+f9T1iT9GByDxfZKAQTCR3kQA=
 golang.org/x/exp v0.0.0-20240719175910-8a7402abbf56 h1:2dVuKD2vS7b0QIHQbpyTISPd0LeHDbnYEryqj5Q1ug8=
 golang.org/x/exp v0.0.0-20240719175910-8a7402abbf56/go.mod h1:M4RDyNAINzryxdtnbRXRL/OHtkFuWGRjvuhBJpk2IlY=
@@ -164,6 +170,8 @@ golang.org/x/lint v0.0.0-20190227174305-5b3e6a55c961/go.mod h1:wehouNa3lNwaWXcvx
 golang.org/x/lint v0.0.0-20190313153728-d0100b6bd8b3/go.mod h1:6SW0HCj/g11FgYtHlgUYUwCkIfeOF89ocIRzGO/8vkc=
 golang.org/x/mod v0.17.0 h1:zY54UmvipHiNd+pm+m0x9KhZ9hl1/7QNMyxXbc6ICqA=
 golang.org/x/mod v0.17.0/go.mod h1:hTbmBsO62+eylJbnUtE2MGJUyE7QWk4xUqPFrRgJ+7c=
+golang.org/x/mod v0.22.0 h1:D4nJWe9zXqHOmWqj4VMOJhvzj7bEZg4wEYa759z1pH4=
+golang.org/x/mod v0.22.0/go.mod h1:6SkKJ3Xj0I0BrPOZoBy3bdMptDDU9oJrpohJ3eWZ1fY=
 golang.org/x/net v0.0.0-20180724234803-3673e40ba225/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4=
 golang.org/x/net v0.0.0-20180826012351-8a410e7b638d/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4=
 golang.org/x/net v0.0.0-20190213061140-3a22650c66bd/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4=
@@ -173,6 +181,8 @@ golang.org/x/net v0.0.0-20201110031124-69a78807bb2b/go.mod h1:sp8m0HH+o8qH0wwXwY
 golang.org/x/net v0.0.0-20211112202133-69e39bad7dc2/go.mod h1:9nx3DQGgdP8bBQD5qxJ1jj9UTztislL4KSBs9R2vV5Y=
 golang.org/x/net v0.33.0 h1:74SYHlV8BIgHIFC/LrYkOGIwL19eTYXQ5wc6TBuO36I=
 golang.org/x/net v0.33.0/go.mod h1:HXLR5J+9DxmrqMwG9qjGCxZ+zKXxBru04zlTvWlWuN4=
+golang.org/x/net v0.34.0 h1:Mb7Mrk043xzHgnRM88suvJFwzVrRfHEHJEl5/71CKw0=
+golang.org/x/net v0.34.0/go.mod h1:di0qlW3YNM5oh6GqDGQr92MyTozJPmybPK4Ev/Gm31k=
 golang.org/x/oauth2 v0.0.0-20180821212333-d2e6202438be/go.mod h1:N/0e6XlmueqKjAGxoOufVs8QHGRruUQn6yWY3a++T0U=
 golang.org/x/oauth2 v0.25.0 h1:CY4y7XT9v0cRI9oupztF8AgiIu99L/ksR/Xp/6jrZ70=
 golang.org/x/oauth2 v0.25.0/go.mod h1:XYTD2NtWslqkgxebSiOHnXEap4TF09sJSc7H1sXbhtI=
@@ -193,23 +203,32 @@ golang.org/x/sys v0.0.0-20210615035016-665e8c7367d1/go.mod h1:oPkhp1MJrh7nUepCBc
 golang.org/x/sys v0.0.0-20220715151400-c0bba94af5f8/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.28.0 h1:Fksou7UEQUWlKvIdsqzJmUmCX3cZuD2+P3XyyzwMhlA=
 golang.org/x/sys v0.28.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
+golang.org/x/sys v0.29.0 h1:TPYlXGxvx1MGTn2GiZDhnjPA9wZzZeGKHHmKhHYvgaU=
+golang.org/x/sys v0.29.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
 golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo=
 golang.org/x/term v0.27.0 h1:WP60Sv1nlK1T6SupCHbXzSaN0b9wUmsPoRS9b61A23Q=
 golang.org/x/term v0.27.0/go.mod h1:iMsnZpn0cago0GOrHO2+Y7u7JPn5AylBrcoWkElMTSM=
+golang.org/x/term v0.28.0 h1:/Ts8HFuMR2E6IP/jlo7QVLZHggjKQbhu/7H0LJFr3Gg=
 golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
 golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
 golang.org/x/text v0.3.6/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
 golang.org/x/text v0.21.0 h1:zyQAAkrwaneQ066sspRyJaG9VNi/YJ1NfzcGB3hZ/qo=
 golang.org/x/text v0.21.0/go.mod h1:4IBbMaMmOPCJ8SecivzSH54+73PCFmPWxNTLm+vZkEQ=
 golang.org/x/time v0.5.0 h1:o7cqy6amK/52YcAKIPlM3a+Fpj35zvRj2TP+e1xFSfk=
 golang.org/x/time v0.5.0/go.mod h1:3BpzKBy/shNhVucY/MWOyx10tF3SFh9QdLuxbVysPQM=
+golang.org/x/time v0.9.0 h1:EsRrnYcQiGH+5FfbgvV4AP7qEZstoyrHB0DzarOQ4ZY=
+golang.org/x/time v0.9.0/go.mod h1:3BpzKBy/shNhVucY/MWOyx10tF3SFh9QdLuxbVysPQM=
 golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
 golang.org/x/tools v0.0.0-20190114222345-bf090417da8b/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
 golang.org/x/tools v0.0.0-20190226205152-f727befe758c/go.mod h1:9Yl7xja0Znq3iFh3HoIrodX9oNMXvdceNzlUR8zjMvY=
 golang.org/x/tools v0.0.0-20190311212946-11955173bddd/go.mod h1:LCzVGOaR6xXOjkQ3onu1FJEFr0SW1gC7cKk1uF8kGRs=
 golang.org/x/tools v0.0.0-20190524140312-2c0ae7006135/go.mod h1:RgjU9mgBXZiqYHBnxXauZ1Gv1EHHAz9KjViQ78xBX0Q=
 golang.org/x/tools v0.21.1-0.20240508182429-e35e4ccd0d2d h1:vU5i/LfpvrRCpgM/VPfJLg5KjxD3E+hfT1SH+d9zLwg=
 golang.org/x/tools v0.21.1-0.20240508182429-e35e4ccd0d2d/go.mod h1:aiJjzUbINMkxbQROHiO6hDPo2LHcIPhhQsa9DLh0yGk=
+golang.org/x/tools v0.29.0 h1:Xx0h3TtM9rzQpQuR4dKLrdglAmCEN5Oi+P74JdhdzXE=
+golang.org/x/tools v0.29.0/go.mod h1:KMQVMRsVxU6nHCFXrBPhDB8XncLNLM0lIy/F14RP588=
+golang.org/x/vulndb v0.0.0-20250130185858-953c816d6fd6 h1:ec2Gsr3f0erz7CL0kMbvkKf6dpRVw7UYe5oygDGEto0=
+golang.org/x/vulndb v0.0.0-20250130185858-953c816d6fd6/go.mod h1:OOULUo0zp2ZeTz9C2YGdkof7jyBG2oSM8Il0VythBvM=
 golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
 google.golang.org/api v0.186.0 h1:n2OPp+PPXX0Axh4GuSsL5QL8xQCTb2oDwyzPnQvqUug=
 google.golang.org/api v0.186.0/go.mod h1:hvRbBmgoje49RV3xqVXrmP6w93n6ehGgIVPYrGtBFFc=

main.go

@@ -3,12 +3,6 @@ package main
 import (
 	"context"
 	"fmt"
-	"log"
-	"os"
-	"strings"
-	"time"
-	"math/rand"
-	"strconv"
 	"github.com/go-git/go-git/v5"
 	"github.com/go-git/go-git/v5/config"
 	"github.com/go-git/go-git/v5/plumbing"
@@ -17,6 +11,12 @@ import (
 	"github.com/google/go-github/v57/github"
 	"github.com/joho/godotenv"
 	"golang.org/x/oauth2"
+	"log"
+	"math/rand"
+	"os"
+	"strconv"
+	"strings"
+	"time"
 )
 
 func main() {
@@ -41,6 +41,8 @@ func main() {
 	prTitle := os.Getenv("PR_TITLE")
 	prBody := os.Getenv("PR_BODY")
 
+	_, pythonPackage, pythonVersion := getPythonVuln(githubPersonalAccessToken)
+
 	bugId := strconv.Itoa(rand.Intn(3000))
 
 	if commitMsg == "<auto>" {
@@ -49,20 +51,19 @@ func main() {
 	if prTitle == "<auto>" {
 		prTitle = generateCommit("Generate a concise pull request title with no placeholders for a GitHub pull request related to the commit" + commitMsg)
 	}
-	if prBody== "<auto>" {
+	if prBody == "<auto>" {
 		prBody = generateCommit("Write a GitHub pull request body for a pull request with the title '" + prTitle + "'.  Make the request body complete with no placeholders or sample text")
 	}
 	if branchPrefix == "<auto>" {
 		branchPrefix = generateCommit("Generate a GitHub branch name with no markdown to hold the pull request " + prTitle + " for bug ID " + bugId)
 		branchPrefix = strings.TrimSuffix(branchPrefix, "\n")
 	}
 
-
 	fmt.Println(commitMsg)
 	fmt.Println(prTitle)
 	fmt.Println(prBody)
 	fmt.Println(branchPrefix)
-	
+
 	//Get code
 	branchName, repo, auth, err := GetCode(branchPrefix, repoURL, githubPersonalAccessToken)
 	if err != nil {
@@ -107,6 +108,9 @@ func main() {
 		return
 	}
 
+	// Create a Python vuln if requirements.txt exists
+	writePythonVuln(pythonPackage, pythonVersion)
+
 	// Stage changes
 	fmt.Println("Staging changes...")
 	_, err = w.Add(".")