PluginVulnerability.go

package wpfinger

import (
	"encoding/json"
	"errors"
	"github.com/hashicorp/go-version"
	"github.com/k0kubun/go-ansi"
	"github.com/schollz/progressbar/v3"
	"gorm.io/gorm"
	"net/http"
	"strings"
)

type WordfenceResponse map[string]WordfenceVulnerability

type WordfenceVulnerability struct {
	Id       string `json:"id"`
	Software []struct {
		Type             string `json:"type"`
		Name             string `json:"name"`
		Slug             string `json:"slug"`
		AffectedVersions map[string]struct {
			FromVersion   string `json:"from_version"`
			FromInclusive bool   `json:"from_inclusive"`
			ToVersion     string `json:"to_version"`
			ToInclusive   bool   `json:"to_inclusive"`
		} `json:"affected_versions"`
	} `json:"software"`
	CVSS struct {
		Vector string  `json:"vector"`
		Score  float64 `json:"score"`
		Rating string  `json:"rating"`
	}
	CWE struct {
		Id          int    `json:"id"`
		Name        string `json:"name"`
		Description string `json:"description"`
	}
	CVE string
}

func UpdateWordfenceDB(db *gorm.DB) {
	setupPogressBar()
	req, err := http.NewRequest(http.MethodGet, "https://www.wordfence.com/api/intelligence/v2/vulnerabilities/production", nil)
	if err != nil {
		panic(err)
	}
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		panic(err)
	}
	if resp.StatusCode != 200 {
		panic(errors.New("not found"))
	}
	var wfResponse WordfenceResponse
	decoder := json.NewDecoder(resp.Body)
	err = decoder.Decode(&wfResponse)

	if err != nil {
		panic(err)
	}
	for _, vulnerability := range wfResponse {
		vulnBar.ChangeMax(len(wfResponse))
		vulnBar.Describe("[cyan][2/2][reset] Saving vulnerabilities in database")
		for _, software := range vulnerability.Software {
			for affectedVersionLabel, affectedVersion := range software.AffectedVersions {
				affectedVersion.FromVersion = strings.ReplaceAll(affectedVersion.FromVersion, "*", "0.0.0")
				affectedVersion.FromVersion = strings.ReplaceAll(affectedVersion.FromVersion, "*", "999999.0.0")
				vulnId := vulnerability.CVE
				if len(vulnId) < 1 {
					continue
				}
				vuln := Vulnerability{
					Slug:            software.Slug,
					AffectedVersion: affectedVersionLabel,
					Type:            software.Type,
					Id:              vulnerability.CVE,
					FromVersion:     affectedVersion.FromVersion,
					FromInclusive:   affectedVersion.FromInclusive,
					ToVersion:       affectedVersion.ToVersion,
					ToInclusive:     affectedVersion.ToInclusive,
					Severity:        strings.ToLower(vulnerability.CVSS.Rating),
				}
				tx := db.Save(&vuln)
				if tx.Error != nil {
					panic(err)
				}
			}
		}
		vulnBar.Add(1)
	}
}

var vulnBar *progressbar.ProgressBar

func setupPogressBar() {
	vulnBar = progressbar.NewOptions(1,
		progressbar.OptionClearOnFinish(),
		progressbar.OptionSetWriter(ansi.NewAnsiStderr()),
		progressbar.OptionEnableColorCodes(true),
		progressbar.OptionShowIts(),
		progressbar.OptionSetDescription("[cyan][1/2][reset] Downloading database from Wordfence..."),
		progressbar.OptionShowDescriptionAtLineEnd(),
		progressbar.OptionSetTheme(progressbar.Theme{
			Saucer:        "[green]=[reset]",
			SaucerHead:    "[green]>[reset]",
			SaucerPadding: " ",
			BarStart:      "[",
			BarEnd:        "]",
		}))
	vulnBar.RenderBlank()
}

type Vulnerability struct {
	Slug            string `gorm:"primaryKey,index"`
	AffectedVersion string `gorm:"primaryKey"`
	Type            string `gorm:"primaryKey,index"`
	Id              string
	FromVersion     string
	FromInclusive   bool
	ToVersion       string
	ToInclusive     bool
	Severity        string
}

func (v Vulnerability) Check(checkVersion string) (bool, error) {
	pVersion, err := version.NewVersion(checkVersion)
	if err != nil {
		return false, err
	}
	fromVersion, err := version.NewVersion(v.FromVersion)
	if err != nil {
		return false, err
	}
	toVersion, err := version.NewVersion(v.ToVersion)
	if err != nil {
		return false, err
	}
	if !v.FromInclusive && pVersion.LessThan(fromVersion) {
		return false, nil
	}
	if v.FromInclusive && pVersion.LessThanOrEqual(fromVersion) {
		return false, nil
	}
	if v.ToInclusive && pVersion.GreaterThan(toVersion) {
		return false, nil
	}
	if !v.ToInclusive && pVersion.GreaterThanOrEqual(toVersion) {
		return false, nil
	}
	return true, nil
}