adding automatic service check on query (projectdiscovery#5291)

* adding automatic service check on query * automatic service check
Ralaam · Jul 8, 2024 · 1c24ced · 1c24ced
1 parent 5cb32a4
commit 1c24ced
Show file tree

Hide file tree

Showing 2 changed files with 58 additions and 4 deletions.
diff --git a/pkg/js/libs/mysql/mysql.go b/pkg/js/libs/mysql/mysql.go
@@ -36,6 +36,7 @@ type (
 // const isMySQL = mysql.IsMySQL('acme.com', 3306);
 // ```
 func (c *MySQLClient) IsMySQL(host string, port int) (bool, error) {
+	// todo: why this is exposed? Service fingerprint should be automatic
 	return memoizedisMySQL(host, port)
 }
 
@@ -77,6 +78,16 @@ func (c *MySQLClient) Connect(host string, port int, username, password string)
 		// host is not valid according to network policy
 		return false, protocolstate.ErrHostDenied.Msgf(host)
 	}
+
+	// executing queries implies the remote mysql service
+	ok, err := c.IsMySQL(host, port)
+	if err != nil {
+		return false, err
+	}
+	if !ok {
+		return false, fmt.Errorf("not a mysql service")
+	}
+
 	dsn, err := BuildDSN(MySQLOptions{
 		Host:     host,
 		Port:     port,
@@ -182,6 +193,16 @@ func (c *MySQLClient) ExecuteQueryWithOpts(opts MySQLOptions, query string) (*ut
 		// host is not valid according to network policy
 		return nil, protocolstate.ErrHostDenied.Msgf(opts.Host)
 	}
+
+	// executing queries implies the remote mysql service
+	ok, err := c.IsMySQL(opts.Host, opts.Port)
+	if err != nil {
+		return nil, err
+	}
+	if !ok {
+		return nil, fmt.Errorf("not a mysql service")
+	}
+
 	dsn, err := BuildDSN(opts)
 	if err != nil {
 		return nil, err
@@ -220,6 +241,15 @@ func (c *MySQLClient) ExecuteQueryWithOpts(opts MySQLOptions, query string) (*ut
 // log(to_json(result));
 // ```
 func (c *MySQLClient) ExecuteQuery(host string, port int, username, password, query string) (*utils.SQLResult, error) {
+	// executing queries implies the remote mysql service
+	ok, err := c.IsMySQL(host, port)
+	if err != nil {
+		return nil, err
+	}
+	if !ok {
+		return nil, fmt.Errorf("not a mysql service")
+	}
+
 	return c.ExecuteQueryWithOpts(MySQLOptions{
 		Host:     host,
 		Port:     port,

diff --git a/pkg/js/libs/postgres/postgres.go b/pkg/js/libs/postgres/postgres.go
@@ -37,6 +37,7 @@ type (
 // const isPostgres = postgres.IsPostgres('acme.com', 5432);
 // ```
 func (c *PGClient) IsPostgres(host string, port int) (bool, error) {
+	// todo: why this is exposed? Service fingerprint should be automatic
 	return memoizedisPostgres(host, port)
 }
 
@@ -74,6 +75,13 @@ func 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) {
+	ok, err := c.IsPostgres(host, port)
+	if err != nil {
+		return false, err
+	}
+	if !ok {
+		return false, fmt.Errorf("not a postgres service")
+	}
 	return memoizedconnect(host, port, username, password, "postgres")
 }
 
@@ -88,6 +96,14 @@ 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) {
+	ok, err := c.IsPostgres(host, port)
+	if err != nil {
+		return nil, err
+	}
+	if !ok {
+		return nil, fmt.Errorf("not a postgres service")
+	}
+
 	return memoizedexecuteQuery(host, port, username, password, dbName, query)
 }
 
@@ -129,6 +145,14 @@ func executeQuery(host string, port int, username string, password string, 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) {
+	ok, err := c.IsPostgres(host, port)
+	if err != nil {
+		return false, err
+	}
+	if !ok {
+		return false, fmt.Errorf("not a postgres service")
+	}
+
 	return memoizedconnect(host, port, username, password, dbName)
 }
 
@@ -149,10 +173,10 @@ func connect(host string, port int, username string, password string, dbName str
 	defer cancel()
 
 	db := pg.Connect(&pg.Options{
-		Addr:               target,
-		User:               username,
-		Password:           password,
-		Database:           dbName,
+		Addr:     target,
+		User:     username,
+		Password: password,
+		Database: dbName,
 		Dialer: func(network, addr string) (net.Conn, error) {
 			return protocolstate.Dialer.Dial(context.Background(), network, addr)
 		},