diff --git a/Makefile b/Makefile index 1ddf950e..cd2c2157 100644 --- a/Makefile +++ b/Makefile @@ -42,4 +42,7 @@ jsupdate: ts: $(GOBUILD) $(GOFLAGS) -ldflags '$(LDFLAGS)' -o "tsgen" pkg/js/devtools/tsgen/cmd/tsgen/main.go ./tsgen -dir pkg/js/libs -out pkg/js/generated/ts +memogen: + $(GOBUILD) $(GOFLAGS) -ldflags '$(LDFLAGS)' -o "memogen" cmd/memogen/memogen.go + ./memogen -src pkg/js/libs -tpl cmd/memogen/function.tpl diff --git a/cmd/memogen/function.tpl b/cmd/memogen/function.tpl new file mode 100644 index 00000000..0b6155b7 --- /dev/null +++ b/cmd/memogen/function.tpl @@ -0,0 +1,28 @@ +// Warning - This is generated code +package {{.SourcePackage}} + +import ( + "github.com/projectdiscovery/utils/memoize" + + {{range .Imports}} + {{.Name}} {{.Path}} + {{end}} +) + +{{range .Functions}} + {{ .SignatureWithPrefix "memoized" }} { + hash := "{{ .Name }}" {{range .Params}} + ":" + fmt.Sprint({{.Name}}) {{end}} + + v, err, _ := protocolstate.Memoizer.Do(hash, func() (interface{}, error) { + return {{.Name}}({{.ParamsNames}}) + }) + if err != nil { + return {{.ResultFirstFieldDefaultValue}}, err + } + if value, ok := v.({{.ResultFirstFieldType}}); ok { + return value, nil + } + + return {{.ResultFirstFieldDefaultValue}}, errors.New("could not convert cached result") + } +{{end}} \ No newline at end of file diff --git a/cmd/memogen/memogen.go b/cmd/memogen/memogen.go new file mode 100644 index 00000000..b494235a --- /dev/null +++ b/cmd/memogen/memogen.go @@ -0,0 +1,77 @@ +// this small cli tool is specific for those functions with arbitrary parameters and with result-error tuple as return values +// func(x,y) => result, error +// it works by creating a new memoized version of the functions in the same path as memo.original.file.go +// some parts are specific for nuclei and hardcoded within the template +package main + +import ( + "flag" + "io/fs" + "log" + "os" + "path/filepath" + + "github.com/projectdiscovery/utils/memoize" + stringsutil "github.com/projectdiscovery/utils/strings" +) + +var ( + srcPath = flag.String("src", "", "nuclei source path") + tplPath = flag.String("tpl", "function.tpl", "template path") + tplSrc []byte +) + +func main() { + flag.Parse() + + var err error + tplSrc, err = os.ReadFile(*tplPath) + if err != nil { + log.Fatal(err) + } + + err = filepath.Walk(*srcPath, walk) + if err != nil { + log.Fatal(err) + } +} + +func walk(path string, info fs.FileInfo, err error) error { + if info.IsDir() { + return nil + } + + if err != nil { + return err + } + + ext := filepath.Ext(path) + base := filepath.Base(path) + + if !stringsutil.EqualFoldAny(ext, ".go") { + return nil + } + + basePath := filepath.Dir(path) + outPath := filepath.Join(basePath, "memo."+base) + + // filename := filepath.Base(path) + data, err := os.ReadFile(path) + if err != nil { + return err + } + if !stringsutil.ContainsAnyI(string(data), "@memo") { + return nil + } + log.Println("processing:", path) + out, err := memoize.Src(string(tplSrc), path, data, "") + if err != nil { + return err + } + + if err := os.WriteFile(outPath, out, os.ModePerm); err != nil { + return err + } + + return nil +} diff --git a/go.mod b/go.mod index be701290..5356486d 100644 --- a/go.mod +++ b/go.mod @@ -91,7 +91,7 @@ require ( github.com/projectdiscovery/tlsx v1.1.6 github.com/projectdiscovery/uncover v1.0.7 github.com/projectdiscovery/useragent v0.0.39 - github.com/projectdiscovery/utils v0.0.80 + github.com/projectdiscovery/utils v0.0.81 github.com/projectdiscovery/wappalyzergo v0.0.111 github.com/redis/go-redis/v9 v9.1.0 github.com/sashabaranov/go-openai v1.15.3 @@ -123,6 +123,7 @@ require ( github.com/bits-and-blooms/bloom/v3 v3.5.0 // indirect github.com/bytedance/sonic v1.9.1 // indirect github.com/cenkalti/backoff/v4 v4.2.1 // indirect + github.com/cespare/xxhash v1.1.0 // indirect github.com/cespare/xxhash/v2 v2.2.0 // indirect github.com/cheggaaa/pb/v3 v3.1.4 // indirect github.com/chenzhuoyu/base64x v0.0.0-20221115062448-fe3a3abad311 // indirect @@ -132,7 +133,6 @@ require ( github.com/corpix/uarand v0.2.0 // indirect github.com/cyphar/filepath-securejoin v0.2.4 // indirect github.com/davidmz/go-pageant v1.0.2 // indirect - github.com/denisbrodbeck/machineid v1.0.1 // indirect github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f // indirect github.com/dlclark/regexp2 v1.10.0 // indirect github.com/docker/cli v24.0.5+incompatible // indirect @@ -188,6 +188,7 @@ require ( github.com/projectdiscovery/asnmap v1.0.6 // indirect github.com/projectdiscovery/cdncheck v1.0.9 // indirect github.com/projectdiscovery/freeport v0.0.5 // indirect + github.com/projectdiscovery/machineid v0.0.0-20240226150047-2e2c51e35983 // indirect github.com/projectdiscovery/stringsutil v0.0.2 // indirect github.com/quic-go/quic-go v0.40.1 // indirect github.com/refraction-networking/utls v1.6.1 // indirect diff --git a/go.sum b/go.sum index dcfa4ae2..4a8cf8b6 100644 --- a/go.sum +++ b/go.sum @@ -78,6 +78,8 @@ github.com/Mzack9999/ldapserver v1.0.2-0.20211229000134-b44a0d6ad0dd h1:RTWs+wEY github.com/Mzack9999/ldapserver v1.0.2-0.20211229000134-b44a0d6ad0dd/go.mod h1:AqtPw7WNT0O69k+AbPKWVGYeW94TqgMW/g+Ppc8AZr4= github.com/Nvveen/Gotty v0.0.0-20120604004816-cd527374f1e5 h1:TngWCqHvy9oXAN6lEVMRuU21PR1EtLVZJmdB18Gu3Rw= github.com/Nvveen/Gotty v0.0.0-20120604004816-cd527374f1e5/go.mod h1:lmUJ/7eu/Q8D7ML55dXQrVaamCz2vxCfdQBasLZfHKk= +github.com/OneOfOne/xxhash v1.2.2 h1:KMrpdQIwFcEqXDklaen+P1axHaj9BSKzvpUUfnHldSE= +github.com/OneOfOne/xxhash v1.2.2/go.mod h1:HSdplMjZKSmBqAxg5vPj2TmRDmfkzw+cTzAElWljhcU= github.com/ProtonMail/go-crypto v0.0.0-20230217124315-7d5c6f04bbb8/go.mod h1:I0gYDMZ6Z5GRU7l58bNFSkPTFN6Yl12dsUlAZ8xy98g= github.com/ProtonMail/go-crypto v0.0.0-20230828082145-3c4c8a2d2371 h1:kkhsdkhsCvIsutKu5zLMgWtgh9YxGCNAw8Ad8hjwfYg= github.com/ProtonMail/go-crypto v0.0.0-20230828082145-3c4c8a2d2371/go.mod h1:EjAoLdwvbIOoOQr3ihjnSoLZRtE8azugULFRteWMNc0= @@ -211,6 +213,8 @@ github.com/cenkalti/backoff v2.2.1+incompatible/go.mod h1:90ReRw6GdpyfrHakVjL/QH github.com/cenkalti/backoff/v4 v4.2.1 h1:y4OZtCnogmCPw98Zjyt5a6+QwPLGkiQsYW5oUqylYbM= github.com/cenkalti/backoff/v4 v4.2.1/go.mod h1:Y3VNntkOUPxTVeUxJ/G5vcM//AlwfmyYozVcomhLiZE= github.com/census-instrumentation/opencensus-proto v0.2.1/go.mod h1:f6KPmirojxKA12rnyqOA5BBL4O983OfeGPqjHWSTneU= +github.com/cespare/xxhash v1.1.0 h1:a6HrQnmkObjyL+Gs60czilIUGqrzKutQD6XZog3p+ko= +github.com/cespare/xxhash v1.1.0/go.mod h1:XrSqR1VqqWfGrhpAt58auRo0WTKS1nRRg3ghfAqPWnc= github.com/cespare/xxhash/v2 v2.1.1/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= github.com/cespare/xxhash/v2 v2.1.2/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= github.com/cespare/xxhash/v2 v2.2.0 h1:DC2CZ1Ep5Y4k3ZQ899DldepgrayRUGE6BBZ/cd9Cj44= @@ -260,8 +264,6 @@ github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/davidmz/go-pageant v1.0.2 h1:bPblRCh5jGU+Uptpz6LgMZGD5hJoOt7otgT454WvHn0= github.com/davidmz/go-pageant v1.0.2/go.mod h1:P2EDDnMqIwG5Rrp05dTRITj9z2zpGcD9efWSkTNKLIE= -github.com/denisbrodbeck/machineid v1.0.1 h1:geKr9qtkB876mXguW2X6TU4ZynleN6ezuMSRhl4D7AQ= -github.com/denisbrodbeck/machineid v1.0.1/go.mod h1:dJUwb7PTidGDeYyUBmXZ2GphQBbjJCrnectwCyxcUSI= github.com/denisenkom/go-mssqldb v0.12.3 h1:pBSGx9Tq67pBOTLmxNuirNTeB8Vjmf886Kx+8Y+8shw= github.com/denisenkom/go-mssqldb v0.12.3/go.mod h1:k0mtMFOnU+AihqFxPMiF05rtiDrorD1Vrm1KEz5hxDo= github.com/dgrijalva/jwt-go v3.2.0+incompatible/go.mod h1:E3ru+11k8xSBh+hMPgOLZmtrrCbhqsmaPHjLKYnJCaQ= @@ -836,6 +838,8 @@ github.com/projectdiscovery/httpx v1.3.9 h1:jDdoGH+5VVU/jI6dnai1DKNw9USPyCcw+tDh github.com/projectdiscovery/httpx v1.3.9/go.mod h1:a/a5X6e2NLnS/+b3buFadGUpZSolnVkMA7KZdpCdg58= github.com/projectdiscovery/interactsh v1.1.8 h1:mDD+f/oo2tV4Z1WyUync0tgYeJyuiS89Un64Gm6Pvgk= github.com/projectdiscovery/interactsh v1.1.8/go.mod h1:E20ywFb7bL01GcOOk+6VZF48XZ8AZvYvBpULoBUSTbg= +github.com/projectdiscovery/machineid v0.0.0-20240226150047-2e2c51e35983 h1:ZScLodGSezQVwsQDtBSMFp72WDq0nNN+KE/5DHKY5QE= +github.com/projectdiscovery/machineid v0.0.0-20240226150047-2e2c51e35983/go.mod h1:3G3BRKui7nMuDFAZKR/M2hiOLtaOmyukT20g88qRQjI= github.com/projectdiscovery/mapcidr v1.1.16 h1:rjj1w5D6hbTsUQXYClLcGdfBEy9bryclgi70t0vBggo= github.com/projectdiscovery/mapcidr v1.1.16/go.mod h1:rGqpBhStdwOQ2uS62QM9qPsybwMwIhT7CTd2bxoHs8Q= github.com/projectdiscovery/n3iwf v0.0.0-20230523120440-b8cd232ff1f5 h1:L/e8z8yw1pfT6bg35NiN7yd1XKtJap5Nk6lMwQ0RNi8= @@ -862,8 +866,8 @@ github.com/projectdiscovery/uncover v1.0.7 h1:ut+2lTuvmftmveqF5RTjMWAgyLj8ltPQC7 github.com/projectdiscovery/uncover v1.0.7/go.mod h1:HFXgm1sRPuoN0D4oATljPIdmbo/EEh1wVuxQqo/dwFE= github.com/projectdiscovery/useragent v0.0.39 h1:s2jyXdtjVo0MfYYkifx7irrOIoA0JhzhZaBkpcoWgV4= github.com/projectdiscovery/useragent v0.0.39/go.mod h1:wO6GQImJ2IQ5K+GDggS/Rhg6IV9Z2Du6NbqC/um0g0w= -github.com/projectdiscovery/utils v0.0.80 h1:daFuQwhVRtQ14JZs3DnI9ubaX273S8V1dZ+x/vr+YbI= -github.com/projectdiscovery/utils v0.0.80/go.mod h1:WXm3MIzKhgqUtTMwxDIW5bWe5nWkCYqRlZeqin0FqTc= +github.com/projectdiscovery/utils v0.0.81 h1:Cqz6uFncCKWRLqpVHWlnHXaRE3whzH32yZJa/1zOEzU= +github.com/projectdiscovery/utils v0.0.81/go.mod h1:pTGvF08EXa07e2OM+tu8IcnxTeAT34bzAhSW/Efcens= github.com/projectdiscovery/wappalyzergo v0.0.111 h1:A1fLEycJ1zwvIVh3jCL1n5OaKn7KP+pHAsFqQYMXPZM= github.com/projectdiscovery/wappalyzergo v0.0.111/go.mod h1:hc/o+fgM8KtdpFesjfBTmHTwsR+yBd+4kYZW/DGy/x8= github.com/projectdiscovery/yamldoc-go v1.0.4 h1:eZoESapnMw6WAHiVgRwNqvbJEfNHEH148uthhFbG5jE= @@ -959,6 +963,7 @@ github.com/smartystreets/goconvey v1.6.4 h1:fv0U8FUIMPNf1L9lnHLvLhgicrIVChEkdzIK github.com/smartystreets/goconvey v1.6.4/go.mod h1:syvi0/a8iFYH4r/RixwvyeAJjdLS9QV7WQ/tjFTllLA= github.com/soheilhy/cmux v0.1.4/go.mod h1:IM3LyeVVIOuxMH7sFAkER9+bJ4dT7Ms6E4xg4kGIyLM= github.com/sony/gobreaker v0.4.1/go.mod h1:ZKptC7FHNvhBz7dN2LGjPVBz2sZJmc0/PkyDJOjmxWY= +github.com/spaolacci/murmur3 v0.0.0-20180118202830-f09979ecbc72/go.mod h1:JwIasOWyU6f++ZhiEuf87xNszmSA2myDM2Kzu9HwQUA= github.com/spaolacci/murmur3 v1.1.0 h1:7c1g84S4BPRrfL5Xrdp6fOJ206sU9y293DDHaoy0bLI= github.com/spaolacci/murmur3 v1.1.0/go.mod h1:JwIasOWyU6f++ZhiEuf87xNszmSA2myDM2Kzu9HwQUA= github.com/spf13/cast v1.5.1 h1:R+kOtfhWQE6TVQzY+4D7wJLBgkdVasCEFxSUBYBYIlA= diff --git a/memogen b/memogen new file mode 100755 index 00000000..baab3a72 Binary files /dev/null and b/memogen differ diff --git a/pkg/js/libs/mssql/memo.mssql.go b/pkg/js/libs/mssql/memo.mssql.go new file mode 100755 index 00000000..eb169761 --- /dev/null +++ b/pkg/js/libs/mssql/memo.mssql.go @@ -0,0 +1,43 @@ +// Warning - This is generated code +package mssql + +import ( + "errors" + "fmt" + + _ "github.com/denisenkom/go-mssqldb" + + "github.com/projectdiscovery/nuclei/v3/pkg/protocols/common/protocolstate" +) + +func memoizedconnect(host string, port int, username string, password string, dbName string) (bool, error) { + hash := "connect" + ":" + fmt.Sprint(host) + ":" + fmt.Sprint(port) + ":" + fmt.Sprint(username) + ":" + fmt.Sprint(password) + ":" + fmt.Sprint(dbName) + + v, err, _ := protocolstate.Memoizer.Do(hash, func() (interface{}, error) { + return connect(host, port, username, password, dbName) + }) + if err != nil { + return false, err + } + if value, ok := v.(bool); ok { + return value, nil + } + + return false, errors.New("could not convert cached result") +} + +func memoizedisMssql(host string, port int) (bool, error) { + hash := "isMssql" + ":" + fmt.Sprint(host) + ":" + fmt.Sprint(port) + + v, err, _ := protocolstate.Memoizer.Do(hash, func() (interface{}, error) { + return isMssql(host, port) + }) + if err != nil { + return false, err + } + if value, ok := v.(bool); ok { + return value, nil + } + + return false, errors.New("could not convert cached result") +} diff --git a/pkg/js/libs/mssql/mssql.go b/pkg/js/libs/mssql/mssql.go index 3b5c2630..1b0efa42 100644 --- a/pkg/js/libs/mssql/mssql.go +++ b/pkg/js/libs/mssql/mssql.go @@ -36,7 +36,7 @@ type ( // const connected = client.Connect('acme.com', 1433, 'username', 'password'); // ``` func (c *MSSQLClient) Connect(host string, port int, username, password string) (bool, error) { - return connect(host, port, username, password, "master") + return memoizedconnect(host, port, username, password, "master") } // ConnectWithDB connects to MS SQL database using given credentials and database name. @@ -50,10 +50,11 @@ func (c *MSSQLClient) Connect(host string, port int, username, password string) // const connected = client.ConnectWithDB('acme.com', 1433, 'username', 'password', 'master'); // ``` func (c *MSSQLClient) ConnectWithDB(host string, port int, username, password, dbName string) (bool, error) { - return connect(host, port, username, password, dbName) + return memoizedconnect(host, port, username, password, dbName) } -func connect(host string, port int, username, password, dbName string) (bool, error) { +// @memo +func connect(host string, port int, username string, password string, dbName string) (bool, error) { if host == "" || port <= 0 { return false, fmt.Errorf("invalid host or port") } @@ -104,6 +105,11 @@ func connect(host string, port int, username, password, dbName string) (bool, er // const isMssql = mssql.IsMssql('acme.com', 1433); // ``` func (c *MSSQLClient) IsMssql(host string, port int) (bool, error) { + return memoizedisMssql(host, port) +} + +// @memo +func isMssql(host string, port int) (bool, error) { if !protocolstate.IsHostAllowed(host) { // host is not valid according to network policy return false, protocolstate.ErrHostDenied.Msgf(host) diff --git a/pkg/js/libs/mysql/memo.mysql.go b/pkg/js/libs/mysql/memo.mysql.go new file mode 100755 index 00000000..60fda434 --- /dev/null +++ b/pkg/js/libs/mysql/memo.mysql.go @@ -0,0 +1,41 @@ +// Warning - This is generated code +package mysql + +import ( + "errors" + "fmt" + + "github.com/projectdiscovery/nuclei/v3/pkg/protocols/common/protocolstate" +) + +func memoizedisMySQL(host string, port int) (bool, error) { + hash := "isMySQL" + ":" + fmt.Sprint(host) + ":" + fmt.Sprint(port) + + v, err, _ := protocolstate.Memoizer.Do(hash, func() (interface{}, error) { + return isMySQL(host, port) + }) + if err != nil { + return false, err + } + if value, ok := v.(bool); ok { + return value, nil + } + + return false, errors.New("could not convert cached result") +} + +func memoizedfingerprintMySQL(host string, port int) (MySQLInfo, error) { + hash := "fingerprintMySQL" + ":" + fmt.Sprint(host) + ":" + fmt.Sprint(port) + + v, err, _ := protocolstate.Memoizer.Do(hash, func() (interface{}, error) { + return fingerprintMySQL(host, port) + }) + if err != nil { + return MySQLInfo{}, err + } + if value, ok := v.(MySQLInfo); ok { + return value, nil + } + + return MySQLInfo{}, errors.New("could not convert cached result") +} diff --git a/pkg/js/libs/mysql/memo.mysql_private.go b/pkg/js/libs/mysql/memo.mysql_private.go new file mode 100755 index 00000000..19d7e81b --- /dev/null +++ b/pkg/js/libs/mysql/memo.mysql_private.go @@ -0,0 +1,25 @@ +// Warning - This is generated code +package mysql + +import ( + "errors" + "fmt" + + "github.com/projectdiscovery/nuclei/v3/pkg/protocols/common/protocolstate" +) + +func memoizedconnectWithDSN(dsn string) (bool, error) { + hash := "connectWithDSN" + ":" + fmt.Sprint(dsn) + + v, err, _ := protocolstate.Memoizer.Do(hash, func() (interface{}, error) { + return connectWithDSN(dsn) + }) + if err != nil { + return false, err + } + if value, ok := v.(bool); ok { + return value, nil + } + + return false, errors.New("could not convert cached result") +} diff --git a/pkg/js/libs/mysql/mysql.go b/pkg/js/libs/mysql/mysql.go index 4c4ce2e4..0fd2f242 100644 --- a/pkg/js/libs/mysql/mysql.go +++ b/pkg/js/libs/mysql/mysql.go @@ -36,6 +36,11 @@ type ( // const isMySQL = mysql.IsMySQL('acme.com', 3306); // ``` func (c *MySQLClient) IsMySQL(host string, port int) (bool, error) { + return memoizedisMySQL(host, port) +} + +// @memo +func isMySQL(host string, port int) (bool, error) { if !protocolstate.IsHostAllowed(host) { // host is not valid according to network policy return false, protocolstate.ErrHostDenied.Msgf(host) @@ -110,6 +115,11 @@ type ( // log(to_json(info)); // ``` func (c *MySQLClient) FingerprintMySQL(host string, port int) (MySQLInfo, error) { + return memoizedfingerprintMySQL(host, port) +} + +// @memo +func fingerprintMySQL(host string, port int) (MySQLInfo, error) { info := MySQLInfo{} if !protocolstate.IsHostAllowed(host) { // host is not valid according to network policy @@ -153,7 +163,7 @@ func (c *MySQLClient) FingerprintMySQL(host string, port int) (MySQLInfo, error) // const connected = client.ConnectWithDSN('username:password@tcp(acme.com:3306)/'); // ``` func (c *MySQLClient) ConnectWithDSN(dsn string) (bool, error) { - return connectWithDSN(dsn) + return memoizedconnectWithDSN(dsn) } // ExecuteQueryWithOpts connects to Mysql database using given credentials diff --git a/pkg/js/libs/mysql/mysql_private.go b/pkg/js/libs/mysql/mysql_private.go index dee0d323..aa89df26 100644 --- a/pkg/js/libs/mysql/mysql_private.go +++ b/pkg/js/libs/mysql/mysql_private.go @@ -65,6 +65,7 @@ func BuildDSN(opts MySQLOptions) (string, error) { return dsn.String(), nil } +// @memo func connectWithDSN(dsn string) (bool, error) { db, err := sql.Open("mysql", dsn) if err != nil { diff --git a/pkg/js/libs/oracle/memo.oracle.go b/pkg/js/libs/oracle/memo.oracle.go new file mode 100755 index 00000000..451f2f64 --- /dev/null +++ b/pkg/js/libs/oracle/memo.oracle.go @@ -0,0 +1,25 @@ +// Warning - This is generated code +package oracle + +import ( + "errors" + "fmt" + + "github.com/projectdiscovery/nuclei/v3/pkg/protocols/common/protocolstate" +) + +func memoizedisOracle(host string, port int) (IsOracleResponse, error) { + hash := "isOracle" + ":" + fmt.Sprint(host) + ":" + fmt.Sprint(port) + + v, err, _ := protocolstate.Memoizer.Do(hash, func() (interface{}, error) { + return isOracle(host, port) + }) + if err != nil { + return IsOracleResponse{}, err + } + if value, ok := v.(IsOracleResponse); ok { + return value, nil + } + + return IsOracleResponse{}, errors.New("could not convert cached result") +} diff --git a/pkg/js/libs/oracle/oracle.go b/pkg/js/libs/oracle/oracle.go index 424c83fa..9e432642 100644 --- a/pkg/js/libs/oracle/oracle.go +++ b/pkg/js/libs/oracle/oracle.go @@ -33,6 +33,11 @@ type ( // log(toJSON(isOracle)); // ``` func IsOracle(host string, port int) (IsOracleResponse, error) { + return memoizedisOracle(host, port) +} + +// @memo +func isOracle(host string, port int) (IsOracleResponse, error) { resp := IsOracleResponse{} timeout := 5 * time.Second diff --git a/pkg/js/libs/pop3/memo.pop3.go b/pkg/js/libs/pop3/memo.pop3.go new file mode 100755 index 00000000..dbd5e463 --- /dev/null +++ b/pkg/js/libs/pop3/memo.pop3.go @@ -0,0 +1,25 @@ +// Warning - This is generated code +package pop3 + +import ( + "errors" + "fmt" + + "github.com/projectdiscovery/nuclei/v3/pkg/protocols/common/protocolstate" +) + +func memoizedisPoP3(host string, port int) (IsPOP3Response, error) { + hash := "isPoP3" + ":" + fmt.Sprint(host) + ":" + fmt.Sprint(port) + + v, err, _ := protocolstate.Memoizer.Do(hash, func() (interface{}, error) { + return isPoP3(host, port) + }) + if err != nil { + return IsPOP3Response{}, err + } + if value, ok := v.(IsPOP3Response); ok { + return value, nil + } + + return IsPOP3Response{}, errors.New("could not convert cached result") +} diff --git a/pkg/js/libs/pop3/pop3.go b/pkg/js/libs/pop3/pop3.go index 165885ba..2662befd 100644 --- a/pkg/js/libs/pop3/pop3.go +++ b/pkg/js/libs/pop3/pop3.go @@ -34,6 +34,11 @@ type ( // log(toJSON(isPOP3)); // ``` func IsPOP3(host string, port int) (IsPOP3Response, error) { + return memoizedisPoP3(host, port) +} + +// @memo +func isPoP3(host string, port int) (IsPOP3Response, error) { resp := IsPOP3Response{} timeout := 5 * time.Second diff --git a/pkg/js/libs/postgres/memo.postgres.go b/pkg/js/libs/postgres/memo.postgres.go new file mode 100755 index 00000000..786d8713 --- /dev/null +++ b/pkg/js/libs/postgres/memo.postgres.go @@ -0,0 +1,61 @@ +// Warning - This is generated code +package postgres + +import ( + "errors" + "fmt" + + _ "github.com/lib/pq" + + utils "github.com/projectdiscovery/nuclei/v3/pkg/js/utils" + + "github.com/projectdiscovery/nuclei/v3/pkg/protocols/common/protocolstate" +) + +func memoizedisPostgres(host string, port int) (bool, error) { + hash := "isPostgres" + ":" + fmt.Sprint(host) + ":" + fmt.Sprint(port) + + v, err, _ := protocolstate.Memoizer.Do(hash, func() (interface{}, error) { + return isPostgres(host, port) + }) + if err != nil { + return false, err + } + if value, ok := v.(bool); ok { + return value, nil + } + + return false, errors.New("could not convert cached result") +} + +func memoizedexecuteQuery(host string, port int, username string, password string, dbName string, query string) (*utils.SQLResult, error) { + hash := "executeQuery" + ":" + fmt.Sprint(host) + ":" + fmt.Sprint(port) + ":" + fmt.Sprint(username) + ":" + fmt.Sprint(password) + ":" + fmt.Sprint(dbName) + ":" + fmt.Sprint(query) + + v, err, _ := protocolstate.Memoizer.Do(hash, func() (interface{}, error) { + return executeQuery(host, port, username, password, dbName, query) + }) + if err != nil { + return nil, err + } + if value, ok := v.(*utils.SQLResult); ok { + return value, nil + } + + return nil, errors.New("could not convert cached result") +} + +func memoizedconnect(host string, port int, username string, password string, dbName string) (bool, error) { + hash := "connect" + ":" + fmt.Sprint(host) + ":" + fmt.Sprint(port) + ":" + fmt.Sprint(username) + ":" + fmt.Sprint(password) + ":" + fmt.Sprint(dbName) + + v, err, _ := protocolstate.Memoizer.Do(hash, func() (interface{}, error) { + return connect(host, port, username, password, dbName) + }) + if err != nil { + return false, err + } + if value, ok := v.(bool); ok { + return value, nil + } + + return false, errors.New("could not convert cached result") +} diff --git a/pkg/js/libs/postgres/postgres.go b/pkg/js/libs/postgres/postgres.go index 3ac0e409..3b52604f 100644 --- a/pkg/js/libs/postgres/postgres.go +++ b/pkg/js/libs/postgres/postgres.go @@ -36,6 +36,11 @@ type ( // const isPostgres = postgres.IsPostgres('acme.com', 5432); // ``` func (c *PGClient) IsPostgres(host string, port int) (bool, error) { + return memoizedisPostgres(host, port) +} + +// @memo +func isPostgres(host string, port int) (bool, error) { timeout := 10 * time.Second conn, err := protocolstate.Dialer.Dial(context.TODO(), "tcp", fmt.Sprintf("%s:%d", host, port)) @@ -68,7 +73,7 @@ func (c *PGClient) IsPostgres(host string, port int) (bool, error) { // const connected = client.Connect('acme.com', 5432, 'username', 'password'); // ``` func (c *PGClient) Connect(host string, port int, username, password string) (bool, error) { - return connect(host, port, username, password, "postgres") + return memoizedconnect(host, port, username, password, "postgres") } // ExecuteQuery connects to Postgres database using given credentials and database name. @@ -82,6 +87,11 @@ func (c *PGClient) Connect(host string, port int, username, password string) (bo // log(to_json(result)); // ``` func (c *PGClient) ExecuteQuery(host string, port int, username, password, dbName, query string) (*utils.SQLResult, error) { + return memoizedexecuteQuery(host, port, username, password, dbName, query) +} + +// @memo +func executeQuery(host string, port int, username string, password string, dbName string, query string) (*utils.SQLResult, error) { if !protocolstate.IsHostAllowed(host) { // host is not valid according to network policy return nil, protocolstate.ErrHostDenied.Msgf(host) @@ -117,10 +127,11 @@ func (c *PGClient) ExecuteQuery(host string, port int, username, password, dbNam // const connected = client.ConnectWithDB('acme.com', 5432, 'username', 'password', 'dbname'); // ``` func (c *PGClient) ConnectWithDB(host string, port int, username, password, dbName string) (bool, error) { - return connect(host, port, username, password, dbName) + return memoizedconnect(host, port, username, password, dbName) } -func connect(host string, port int, username, password, dbName string) (bool, error) { +// @memo +func connect(host string, port int, username string, password string, dbName string) (bool, error) { if host == "" || port <= 0 { return false, fmt.Errorf("invalid host or port") } @@ -138,6 +149,8 @@ func connect(host string, port int, username, password, dbName string) (bool, er Password: password, Database: dbName, }) + defer db.Close() + _, err := db.Exec("select 1") if err != nil { switch true { diff --git a/pkg/js/libs/rdp/memo.rdp.go b/pkg/js/libs/rdp/memo.rdp.go new file mode 100755 index 00000000..c592e20e --- /dev/null +++ b/pkg/js/libs/rdp/memo.rdp.go @@ -0,0 +1,41 @@ +// Warning - This is generated code +package rdp + +import ( + "errors" + "fmt" + + "github.com/projectdiscovery/nuclei/v3/pkg/protocols/common/protocolstate" +) + +func memoizedisRDP(host string, port int) (IsRDPResponse, error) { + hash := "isRDP" + ":" + fmt.Sprint(host) + ":" + fmt.Sprint(port) + + v, err, _ := protocolstate.Memoizer.Do(hash, func() (interface{}, error) { + return isRDP(host, port) + }) + if err != nil { + return IsRDPResponse{}, err + } + if value, ok := v.(IsRDPResponse); ok { + return value, nil + } + + return IsRDPResponse{}, errors.New("could not convert cached result") +} + +func memoizedcheckRDPAuth(host string, port int) (CheckRDPAuthResponse, error) { + hash := "checkRDPAuth" + ":" + fmt.Sprint(host) + ":" + fmt.Sprint(port) + + v, err, _ := protocolstate.Memoizer.Do(hash, func() (interface{}, error) { + return checkRDPAuth(host, port) + }) + if err != nil { + return CheckRDPAuthResponse{}, err + } + if value, ok := v.(CheckRDPAuthResponse); ok { + return value, nil + } + + return CheckRDPAuthResponse{}, errors.New("could not convert cached result") +} diff --git a/pkg/js/libs/rdp/rdp.go b/pkg/js/libs/rdp/rdp.go index d7790e7d..7a985ab2 100644 --- a/pkg/js/libs/rdp/rdp.go +++ b/pkg/js/libs/rdp/rdp.go @@ -36,6 +36,11 @@ type ( // log(toJSON(isRDP)); // ``` func IsRDP(host string, port int) (IsRDPResponse, error) { + return memoizedisRDP(host, port) +} + +// @memo +func isRDP(host string, port int) (IsRDPResponse, error) { resp := IsRDPResponse{} timeout := 5 * time.Second @@ -82,6 +87,11 @@ type ( // log(toJSON(checkRDPAuth)); // ``` func CheckRDPAuth(host string, port int) (CheckRDPAuthResponse, error) { + return memoizedcheckRDPAuth(host, port) +} + +// @memo +func checkRDPAuth(host string, port int) (CheckRDPAuthResponse, error) { resp := CheckRDPAuthResponse{} timeout := 5 * time.Second diff --git a/pkg/js/libs/redis/memo.redis.go b/pkg/js/libs/redis/memo.redis.go new file mode 100755 index 00000000..d53c4489 --- /dev/null +++ b/pkg/js/libs/redis/memo.redis.go @@ -0,0 +1,73 @@ +// Warning - This is generated code +package redis + +import ( + "errors" + "fmt" + + "github.com/projectdiscovery/nuclei/v3/pkg/protocols/common/protocolstate" +) + +func memoizedgetServerInfo(host string, port int) (string, error) { + hash := "getServerInfo" + ":" + fmt.Sprint(host) + ":" + fmt.Sprint(port) + + v, err, _ := protocolstate.Memoizer.Do(hash, func() (interface{}, error) { + return getServerInfo(host, port) + }) + if err != nil { + return "", err + } + if value, ok := v.(string); ok { + return value, nil + } + + return "", errors.New("could not convert cached result") +} + +func memoizedconnect(host string, port int, password string) (bool, error) { + hash := "connect" + ":" + fmt.Sprint(host) + ":" + fmt.Sprint(port) + ":" + fmt.Sprint(password) + + v, err, _ := protocolstate.Memoizer.Do(hash, func() (interface{}, error) { + return connect(host, port, password) + }) + if err != nil { + return false, err + } + if value, ok := v.(bool); ok { + return value, nil + } + + return false, errors.New("could not convert cached result") +} + +func memoizedgetServerInfoAuth(host string, port int, password string) (string, error) { + hash := "getServerInfoAuth" + ":" + fmt.Sprint(host) + ":" + fmt.Sprint(port) + ":" + fmt.Sprint(password) + + v, err, _ := protocolstate.Memoizer.Do(hash, func() (interface{}, error) { + return getServerInfoAuth(host, port, password) + }) + if err != nil { + return "", err + } + if value, ok := v.(string); ok { + return value, nil + } + + return "", errors.New("could not convert cached result") +} + +func memoizedisAuthenticated(host string, port int) (bool, error) { + hash := "isAuthenticated" + ":" + fmt.Sprint(host) + ":" + fmt.Sprint(port) + + v, err, _ := protocolstate.Memoizer.Do(hash, func() (interface{}, error) { + return isAuthenticated(host, port) + }) + if err != nil { + return false, err + } + if value, ok := v.(bool); ok { + return value, nil + } + + return false, errors.New("could not convert cached result") +} diff --git a/pkg/js/libs/redis/redis.go b/pkg/js/libs/redis/redis.go index 04f7145f..3fb4de83 100644 --- a/pkg/js/libs/redis/redis.go +++ b/pkg/js/libs/redis/redis.go @@ -19,6 +19,11 @@ import ( // const info = redis.GetServerInfo('acme.com', 6379); // ``` func GetServerInfo(host string, port int) (string, error) { + return memoizedgetServerInfo(host, port) +} + +// @memo +func getServerInfo(host string, port int) (string, error) { if !protocolstate.IsHostAllowed(host) { // host is not valid according to network policy return "", protocolstate.ErrHostDenied.Msgf(host) @@ -29,6 +34,7 @@ func GetServerInfo(host string, port int) (string, error) { Password: "", // no password set DB: 0, // use default DB }) + defer client.Close() // Ping the Redis server _, err := client.Ping(context.TODO()).Result() @@ -52,6 +58,11 @@ func GetServerInfo(host string, port int) (string, error) { // const connected = redis.Connect('acme.com', 6379, 'password'); // ``` func Connect(host string, port int, password string) (bool, error) { + return memoizedconnect(host, port, password) +} + +// @memo +func connect(host string, port int, password string) (bool, error) { if !protocolstate.IsHostAllowed(host) { // host is not valid according to network policy return false, protocolstate.ErrHostDenied.Msgf(host) @@ -62,6 +73,8 @@ func Connect(host string, port int, password string) (bool, error) { Password: password, // no password set DB: 0, // use default DB }) + defer client.Close() + _, err := client.Ping(context.TODO()).Result() if err != nil { return false, err @@ -82,6 +95,11 @@ func Connect(host string, port int, password string) (bool, error) { // const info = redis.GetServerInfoAuth('acme.com', 6379, 'password'); // ``` func GetServerInfoAuth(host string, port int, password string) (string, error) { + return memoizedgetServerInfoAuth(host, port, password) +} + +// @memo +func getServerInfoAuth(host string, port int, password string) (string, error) { if !protocolstate.IsHostAllowed(host) { // host is not valid according to network policy return "", protocolstate.ErrHostDenied.Msgf(host) @@ -92,6 +110,7 @@ func GetServerInfoAuth(host string, port int, password string) (string, error) { Password: password, // no password set DB: 0, // use default DB }) + defer client.Close() // Ping the Redis server _, err := client.Ping(context.TODO()).Result() @@ -115,6 +134,11 @@ func GetServerInfoAuth(host string, port int, password string) (string, error) { // const isAuthenticated = redis.IsAuthenticated('acme.com', 6379); // ``` func IsAuthenticated(host string, port int) (bool, error) { + return memoizedisAuthenticated(host, port) +} + +// @memo +func isAuthenticated(host string, port int) (bool, error) { plugin := pluginsredis.REDISPlugin{} timeout := 5 * time.Second conn, err := protocolstate.Dialer.Dial(context.TODO(), "tcp", fmt.Sprintf("%s:%d", host, port)) @@ -147,6 +171,7 @@ func RunLuaScript(host string, port int, password string, script string) (interf Password: password, DB: 0, // use default DB }) + defer client.Close() // Ping the Redis server _, err := client.Ping(context.TODO()).Result() diff --git a/pkg/js/libs/rsync/memo.rsync.go b/pkg/js/libs/rsync/memo.rsync.go new file mode 100755 index 00000000..5cb0d029 --- /dev/null +++ b/pkg/js/libs/rsync/memo.rsync.go @@ -0,0 +1,25 @@ +// Warning - This is generated code +package rsync + +import ( + "errors" + "fmt" + + "github.com/projectdiscovery/nuclei/v3/pkg/protocols/common/protocolstate" +) + +func memoizedisRsync(host string, port int) (IsRsyncResponse, error) { + hash := "isRsync" + ":" + fmt.Sprint(host) + ":" + fmt.Sprint(port) + + v, err, _ := protocolstate.Memoizer.Do(hash, func() (interface{}, error) { + return isRsync(host, port) + }) + if err != nil { + return IsRsyncResponse{}, err + } + if value, ok := v.(IsRsyncResponse); ok { + return value, nil + } + + return IsRsyncResponse{}, errors.New("could not convert cached result") +} diff --git a/pkg/js/libs/rsync/rsync.go b/pkg/js/libs/rsync/rsync.go index 41a334dd..c9cf18f3 100644 --- a/pkg/js/libs/rsync/rsync.go +++ b/pkg/js/libs/rsync/rsync.go @@ -34,6 +34,11 @@ type ( // log(toJSON(isRsync)); // ``` func IsRsync(host string, port int) (IsRsyncResponse, error) { + return memoizedisRsync(host, port) +} + +// @memo +func isRsync(host string, port int) (IsRsyncResponse, error) { resp := IsRsyncResponse{} timeout := 5 * time.Second diff --git a/pkg/js/libs/smb/memo.smb.go b/pkg/js/libs/smb/memo.smb.go new file mode 100755 index 00000000..51d6584f --- /dev/null +++ b/pkg/js/libs/smb/memo.smb.go @@ -0,0 +1,43 @@ +// Warning - This is generated code +package smb + +import ( + "errors" + "fmt" + + "github.com/projectdiscovery/nuclei/v3/pkg/protocols/common/protocolstate" + + "github.com/zmap/zgrab2/lib/smb/smb" +) + +func memoizedconnectSMBInfoMode(host string, port int) (*smb.SMBLog, error) { + hash := "connectSMBInfoMode" + ":" + fmt.Sprint(host) + ":" + fmt.Sprint(port) + + v, err, _ := protocolstate.Memoizer.Do(hash, func() (interface{}, error) { + return connectSMBInfoMode(host, port) + }) + if err != nil { + return nil, err + } + if value, ok := v.(*smb.SMBLog); ok { + return value, nil + } + + return nil, errors.New("could not convert cached result") +} + +func memoizedlistShares(host string, port int, user string, password string) ([]string, error) { + hash := "listShares" + ":" + fmt.Sprint(host) + ":" + fmt.Sprint(port) + ":" + fmt.Sprint(user) + ":" + fmt.Sprint(password) + + v, err, _ := protocolstate.Memoizer.Do(hash, func() (interface{}, error) { + return listShares(host, port, user, password) + }) + if err != nil { + return []string{}, err + } + if value, ok := v.([]string); ok { + return value, nil + } + + return []string{}, errors.New("could not convert cached result") +} diff --git a/pkg/js/libs/smb/memo.smb_private.go b/pkg/js/libs/smb/memo.smb_private.go new file mode 100755 index 00000000..fe47d1a2 --- /dev/null +++ b/pkg/js/libs/smb/memo.smb_private.go @@ -0,0 +1,29 @@ +// Warning - This is generated code +package smb + +import ( + "errors" + "fmt" + + "time" + + "github.com/praetorian-inc/fingerprintx/pkg/plugins" + + "github.com/projectdiscovery/nuclei/v3/pkg/protocols/common/protocolstate" +) + +func memoizedcollectSMBv2Metadata(host string, port int, timeout time.Duration) (*plugins.ServiceSMB, error) { + hash := "collectSMBv2Metadata" + ":" + fmt.Sprint(host) + ":" + fmt.Sprint(port) + ":" + fmt.Sprint(timeout) + + v, err, _ := protocolstate.Memoizer.Do(hash, func() (interface{}, error) { + return collectSMBv2Metadata(host, port, timeout) + }) + if err != nil { + return nil, err + } + if value, ok := v.(*plugins.ServiceSMB); ok { + return value, nil + } + + return nil, errors.New("could not convert cached result") +} diff --git a/pkg/js/libs/smb/memo.smbghost.go b/pkg/js/libs/smb/memo.smbghost.go new file mode 100755 index 00000000..25e9d187 --- /dev/null +++ b/pkg/js/libs/smb/memo.smbghost.go @@ -0,0 +1,25 @@ +// Warning - This is generated code +package smb + +import ( + "errors" + "fmt" + + "github.com/projectdiscovery/nuclei/v3/pkg/protocols/common/protocolstate" +) + +func memoizeddetectSMBGhost(host string, port int) (bool, error) { + hash := "detectSMBGhost" + ":" + fmt.Sprint(host) + ":" + fmt.Sprint(port) + + v, err, _ := protocolstate.Memoizer.Do(hash, func() (interface{}, error) { + return detectSMBGhost(host, port) + }) + if err != nil { + return false, err + } + if value, ok := v.(bool); ok { + return value, nil + } + + return false, errors.New("could not convert cached result") +} diff --git a/pkg/js/libs/smb/smb.go b/pkg/js/libs/smb/smb.go index 46937c5d..2d981281 100644 --- a/pkg/js/libs/smb/smb.go +++ b/pkg/js/libs/smb/smb.go @@ -35,6 +35,11 @@ type ( // log(to_json(info)); // ``` func (c *SMBClient) ConnectSMBInfoMode(host string, port int) (*smb.SMBLog, error) { + return memoizedconnectSMBInfoMode(host, port) +} + +// @memo +func connectSMBInfoMode(host string, port int) (*smb.SMBLog, error) { if !protocolstate.IsHostAllowed(host) { // host is not valid according to network policy return nil, protocolstate.ErrHostDenied.Msgf(host) @@ -44,7 +49,7 @@ func (c *SMBClient) ConnectSMBInfoMode(host string, port int) (*smb.SMBLog, erro return nil, err } // try to get SMBv2/v3 info - result, err := c.getSMBInfo(conn, true, false) + result, err := getSMBInfo(conn, true, false) _ = conn.Close() // close regardless of error if err == nil { return result, nil @@ -56,7 +61,7 @@ func (c *SMBClient) ConnectSMBInfoMode(host string, port int) (*smb.SMBLog, erro return nil, err } defer conn.Close() - result, err = c.getSMBInfo(conn, true, true) + result, err = getSMBInfo(conn, true, true) if err != nil { return result, nil } @@ -79,7 +84,7 @@ func (c *SMBClient) ListSMBv2Metadata(host string, port int) (*plugins.ServiceSM // host is not valid according to network policy return nil, protocolstate.ErrHostDenied.Msgf(host) } - return collectSMBv2Metadata(host, port, 5*time.Second) + return memoizedcollectSMBv2Metadata(host, port, 5*time.Second) } // ListShares tries to connect to provided host and port @@ -98,6 +103,11 @@ func (c *SMBClient) ListSMBv2Metadata(host string, port int) (*plugins.ServiceSM // // ``` func (c *SMBClient) ListShares(host string, port int, user, password string) ([]string, error) { + return memoizedlistShares(host, port, user, password) +} + +// @memo +func listShares(host string, port int, user string, password string) ([]string, error) { if !protocolstate.IsHostAllowed(host) { // host is not valid according to network policy return nil, protocolstate.ErrHostDenied.Msgf(host) diff --git a/pkg/js/libs/smb/smb_private.go b/pkg/js/libs/smb/smb_private.go index 6d766f0e..74f076b1 100644 --- a/pkg/js/libs/smb/smb_private.go +++ b/pkg/js/libs/smb/smb_private.go @@ -15,6 +15,7 @@ import ( // ==== private helper functions/methods ==== // collectSMBv2Metadata collects metadata for SMBv2 services. +// @memo func collectSMBv2Metadata(host string, port int, timeout time.Duration) (*plugins.ServiceSMB, error) { if timeout == 0 { timeout = 5 * time.Second @@ -33,7 +34,7 @@ func collectSMBv2Metadata(host string, port int, timeout time.Duration) (*plugin } // getSMBInfo -func (c *SMBClient) getSMBInfo(conn net.Conn, setupSession, v1 bool) (*zgrabsmb.SMBLog, error) { +func getSMBInfo(conn net.Conn, setupSession, v1 bool) (*zgrabsmb.SMBLog, error) { _ = conn.SetDeadline(time.Now().Add(10 * time.Second)) defer func() { _ = conn.SetDeadline(time.Time{}) diff --git a/pkg/js/libs/smb/smbghost.go b/pkg/js/libs/smb/smbghost.go index 8e77a293..6c8b3c27 100644 --- a/pkg/js/libs/smb/smbghost.go +++ b/pkg/js/libs/smb/smbghost.go @@ -26,6 +26,11 @@ const ( // const isSMBGhost = smb.DetectSMBGhost('acme.com', 445); // ``` func (c *SMBClient) DetectSMBGhost(host string, port int) (bool, error) { + return memoizeddetectSMBGhost(host, port) +} + +// @memo +func detectSMBGhost(host string, port int) (bool, error) { if !protocolstate.IsHostAllowed(host) { // host is not valid according to network policy return false, protocolstate.ErrHostDenied.Msgf(host) diff --git a/pkg/js/libs/ssh/memo.ssh.go b/pkg/js/libs/ssh/memo.ssh.go new file mode 100755 index 00000000..0bfed584 --- /dev/null +++ b/pkg/js/libs/ssh/memo.ssh.go @@ -0,0 +1,27 @@ +// Warning - This is generated code +package ssh + +import ( + "errors" + "fmt" + + "github.com/projectdiscovery/nuclei/v3/pkg/protocols/common/protocolstate" + + "github.com/zmap/zgrab2/lib/ssh" +) + +func memoizedconnectSSHInfoMode(opts *connectOptions) (*ssh.HandshakeLog, error) { + hash := "connectSSHInfoMode" + ":" + fmt.Sprint(opts) + + v, err, _ := protocolstate.Memoizer.Do(hash, func() (interface{}, error) { + return connectSSHInfoMode(opts) + }) + if err != nil { + return nil, err + } + if value, ok := v.(*ssh.HandshakeLog); ok { + return value, nil + } + + return nil, errors.New("could not convert cached result") +} diff --git a/pkg/js/libs/ssh/ssh.go b/pkg/js/libs/ssh/ssh.go index 920b8202..f6639a57 100644 --- a/pkg/js/libs/ssh/ssh.go +++ b/pkg/js/libs/ssh/ssh.go @@ -101,7 +101,7 @@ func (c *SSHClient) ConnectWithKey(host string, port int, username, key string) // log(to_json(info)); // ``` func (c *SSHClient) ConnectSSHInfoMode(host string, port int) (*ssh.HandshakeLog, error) { - return connectSSHInfoMode(&connectOptions{ + return memoizedconnectSSHInfoMode(&connectOptions{ Host: host, Port: port, }) @@ -182,6 +182,7 @@ func (c *connectOptions) validate() error { return nil } +// @memo func connectSSHInfoMode(opts *connectOptions) (*ssh.HandshakeLog, error) { if err := opts.validate(); err != nil { return nil, err diff --git a/pkg/js/libs/telnet/memo.telnet.go b/pkg/js/libs/telnet/memo.telnet.go new file mode 100755 index 00000000..0e29a5e7 --- /dev/null +++ b/pkg/js/libs/telnet/memo.telnet.go @@ -0,0 +1,25 @@ +// Warning - This is generated code +package telnet + +import ( + "errors" + "fmt" + + "github.com/projectdiscovery/nuclei/v3/pkg/protocols/common/protocolstate" +) + +func memoizedisTelnet(host string, port int) (IsTelnetResponse, error) { + hash := "isTelnet" + ":" + fmt.Sprint(host) + ":" + fmt.Sprint(port) + + v, err, _ := protocolstate.Memoizer.Do(hash, func() (interface{}, error) { + return isTelnet(host, port) + }) + if err != nil { + return IsTelnetResponse{}, err + } + if value, ok := v.(IsTelnetResponse); ok { + return value, nil + } + + return IsTelnetResponse{}, errors.New("could not convert cached result") +} diff --git a/pkg/js/libs/telnet/telnet.go b/pkg/js/libs/telnet/telnet.go index 02135a95..d7145475 100644 --- a/pkg/js/libs/telnet/telnet.go +++ b/pkg/js/libs/telnet/telnet.go @@ -34,6 +34,11 @@ type ( // log(toJSON(isTelnet)); // ``` func IsTelnet(host string, port int) (IsTelnetResponse, error) { + return memoizedisTelnet(host, port) +} + +// @memo +func isTelnet(host string, port int) (IsTelnetResponse, error) { resp := IsTelnetResponse{} timeout := 5 * time.Second diff --git a/pkg/js/libs/vnc/memo.vnc.go b/pkg/js/libs/vnc/memo.vnc.go new file mode 100755 index 00000000..8e2fd454 --- /dev/null +++ b/pkg/js/libs/vnc/memo.vnc.go @@ -0,0 +1,25 @@ +// Warning - This is generated code +package vnc + +import ( + "errors" + "fmt" + + "github.com/projectdiscovery/nuclei/v3/pkg/protocols/common/protocolstate" +) + +func memoizedisVNC(host string, port int) (IsVNCResponse, error) { + hash := "isVNC" + ":" + fmt.Sprint(host) + ":" + fmt.Sprint(port) + + v, err, _ := protocolstate.Memoizer.Do(hash, func() (interface{}, error) { + return isVNC(host, port) + }) + if err != nil { + return IsVNCResponse{}, err + } + if value, ok := v.(IsVNCResponse); ok { + return value, nil + } + + return IsVNCResponse{}, errors.New("could not convert cached result") +} diff --git a/pkg/js/libs/vnc/vnc.go b/pkg/js/libs/vnc/vnc.go index 17f2c199..c5d4577c 100644 --- a/pkg/js/libs/vnc/vnc.go +++ b/pkg/js/libs/vnc/vnc.go @@ -35,6 +35,11 @@ type ( // log(toJSON(isVNC)); // ``` func IsVNC(host string, port int) (IsVNCResponse, error) { + return memoizedisVNC(host, port) +} + +// @memo +func isVNC(host string, port int) (IsVNCResponse, error) { resp := IsVNCResponse{} timeout := 5 * time.Second diff --git a/pkg/protocols/common/protocolstate/memoizer.go b/pkg/protocols/common/protocolstate/memoizer.go new file mode 100644 index 00000000..812472d1 --- /dev/null +++ b/pkg/protocols/common/protocolstate/memoizer.go @@ -0,0 +1,15 @@ +package protocolstate + +import ( + "github.com/projectdiscovery/utils/memoize" +) + +var Memoizer *memoize.Memoizer + +func init() { + var err error + Memoizer, err = memoize.New(memoize.WithMaxSize(1500)) + if err != nil { + panic(err) + } +}