redis.go

package db

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"strconv"
	"time"

	"github.com/cheggaaa/pb/v3"
	"github.com/go-redis/redis/v8"
	"github.com/inconshreveable/log15"
	"github.com/spf13/viper"
	"golang.org/x/xerrors"

	"github.com/vulsio/go-exploitdb/config"
	"github.com/vulsio/go-exploitdb/models"
)

/**
# Redis Data Structure
- Sets
  ┌───┬────────────────┬──────────────┬───────────────────────────┐
  │NO │      KEY       │   MEMBER     │            PURPOSE        │
  └───┴────────────────┴──────────────┴───────────────────────────┘
  ┌───┬────────────────┬──────────────┬───────────────────────────┐
  │ 1 │ EDB#CVE#$CVEID │ $EXPLOITDBID │ TO GET EXPLOIT FROM CVEID │
  └───┴────────────────┴──────────────┴───────────────────────────┘

- Hash
  ┌───┬──────────────────────┬───────────────┬───────────────┬───────────────────────────────────────────────────┐
  │NO │    KEY               │   FIELD       │  VALUE        │       PURPOSE                                     │
  └───┴──────────────────────┴───────────────┴───────────────┴───────────────────────────────────────────────────┘
  ┌───┬──────────────────────┬───────────────┬───────────────┬───────────────────────────────────────────────────┐
  │ 1 │ EDB#EDB#$EXPLOITDBID │    $CVEID     │ $EXPLOIT JSON │ TO GET EXPLOIT FROM EXPLOITDBID                   │
  ├───┼──────────────────────┼───────────────┼───────────────┼───────────────────────────────────────────────────┤
  │ 2 │ EDB#DEP              │ $EXPLOITTYPE  │     JSON      │ TO DELETE OUTDATED AND UNNEEDED FIELD AND MEMBER  │
  ├───┼──────────────────────┼───────────────┼───────────────┼───────────────────────────────────────────────────┤
  │ 3 │ EDB#FETCHMETA        │   Revision    │    string     │ GET GO-EXPLOIT BINARY REVISION                    │
  ├───┼──────────────────────┼───────────────┼───────────────┼───────────────────────────────────────────────────┤
  │ 4 │ EDB#FETCHMETA        │ SchemaVersion │     uint      │ GET GO-EXPLOIT SCHEMA VERSION                     │
  ├───┼──────────────────────┼───────────────┼───────────────┼───────────────────────────────────────────────────┤
  │ 5 │ EDB#FETCHMETA        │ LastFetchedAt │   time.Time   │ GET GO-EXPLOIT LAST FETCHED TIME                  │
  └───┴──────────────────────┴───────────────┴───────────────┴───────────────────────────────────────────────────┘
**/

const (
	dialectRedis         = "redis"
	depKey               = "EDB#DEP"
	exploitDBIDKeyFormat = "EDB#EDB#%s"
	cveIDKeyFormat       = "EDB#CVE#%s"
	fetchMetaKey         = "EDB#FETCHMETA"
)

// RedisDriver is Driver for Redis
type RedisDriver struct {
	name string
	conn *redis.Client
}

// Name return db name
func (r *RedisDriver) Name() string {
	return r.name
}

// OpenDB opens Database
func (r *RedisDriver) OpenDB(_, dbPath string, _ bool, option Option) error {
	if err := r.connectRedis(dbPath, option); err != nil {
		return xerrors.Errorf("Failed to open DB. dbtype: %s, dbpath: %s, err: %w", dialectRedis, dbPath, err)
	}
	return nil
}

// CloseDB close Database
func (r *RedisDriver) CloseDB() (err error) {
	if r.conn == nil {
		return
	}
	if err = r.conn.Close(); err != nil {
		return xerrors.Errorf("Failed to close DB. Type: %s. err: %w", r.name, err)
	}
	return
}

func (r *RedisDriver) connectRedis(dbPath string, option Option) error {
	opt, err := redis.ParseURL(dbPath)
	if err != nil {
		return xerrors.Errorf("Failed to parse url. err: %w", err)
	}
	if 0 < option.RedisTimeout.Seconds() {
		opt.ReadTimeout = option.RedisTimeout
	}
	r.conn = redis.NewClient(opt)
	return r.conn.Ping(context.Background()).Err()
}

// MigrateDB migrates Database
func (r *RedisDriver) MigrateDB() error {
	return nil
}

// GetExploitByCveID :
func (r *RedisDriver) GetExploitByCveID(cveID string) ([]models.Exploit, error) {
	ctx := context.Background()

	exploitIDs, err := r.conn.SMembers(ctx, fmt.Sprintf(cveIDKeyFormat, cveID)).Result()
	if err != nil {
		return nil, xerrors.Errorf("Failed to SMembers. err: %w", err)
	}
	if len(exploitIDs) == 0 {
		return []models.Exploit{}, nil
	}

	pipe := r.conn.Pipeline()
	for _, exploitID := range exploitIDs {
		_ = pipe.HGet(ctx, fmt.Sprintf(exploitDBIDKeyFormat, exploitID), cveID)
	}
	cmders, err := pipe.Exec(ctx)
	if err != nil {
		return nil, xerrors.Errorf("Failed to exec pipeline. err: %w", err)
	}

	exploits := []models.Exploit{}
	for _, cmder := range cmders {
		str, err := cmder.(*redis.StringCmd).Result()
		if err != nil {
			return nil, xerrors.Errorf("Failed to HGet. err: %w", err)
		}

		var exploit models.Exploit
		if err := json.Unmarshal([]byte(str), &exploit); err != nil {
			return nil, xerrors.Errorf("Failed to Unmarshal json. err: %w", err)
		}
		exploits = append(exploits, exploit)
	}
	return exploits, nil
}

// GetExploitByID :
func (r *RedisDriver) GetExploitByID(exploitDBID string) ([]models.Exploit, error) {
	ctx := context.Background()

	result, err := r.conn.HGetAll(ctx, fmt.Sprintf(exploitDBIDKeyFormat, exploitDBID)).Result()
	if err != nil {
		return nil, xerrors.Errorf("Failed to HGetAll. err: %w", err)
	}

	exploits := []models.Exploit{}
	for _, str := range result {
		var exploit models.Exploit
		if err := json.Unmarshal([]byte(str), &exploit); err != nil {
			return nil, xerrors.Errorf("Failed to Unmarshal json. err: %w", err)
		}
		exploits = append(exploits, exploit)
	}
	return exploits, nil
}

// GetExploitAll :
func (r *RedisDriver) GetExploitAll() ([]models.Exploit, error) {
	ctx := context.Background()

	exploitDBIDkeys := []string{}
	var cursor uint64
	for {
		var keys []string
		var err error
		keys, cursor, err = r.conn.Scan(ctx, cursor, fmt.Sprintf(exploitDBIDKeyFormat, "*"), 10).Result()
		if err != nil {
			return nil, xerrors.Errorf("Failed to Scan. err: %w", err)
		}

		exploitDBIDkeys = append(exploitDBIDkeys, keys...)

		if cursor == 0 {
			break
		}
	}
	if len(exploitDBIDkeys) == 0 {
		return []models.Exploit{}, nil
	}

	pipe := r.conn.Pipeline()
	for _, exploitDBIDKey := range exploitDBIDkeys {
		_ = pipe.HGetAll(ctx, exploitDBIDKey)
	}
	cmders, err := pipe.Exec(ctx)
	if err != nil {
		return nil, xerrors.Errorf("Failed to exec pipeline. err: %w", err)
	}

	exploits := []models.Exploit{}
	for _, cmder := range cmders {
		result, err := cmder.(*redis.StringStringMapCmd).Result()
		if err != nil {
			return nil, xerrors.Errorf("Failed to HGetAll. err: %w", err)
		}
		for _, str := range result {
			var exploit models.Exploit
			if err := json.Unmarshal([]byte(str), &exploit); err != nil {
				return nil, xerrors.Errorf("Failed to Unmarshal json. err: %w", err)
			}

			exploits = append(exploits, exploit)
		}
	}
	return exploits, nil
}

// GetExploitMultiByID :
func (r *RedisDriver) GetExploitMultiByID(exploitUniqueIDs []string) (map[string][]models.Exploit, error) {
	ctx := context.Background()

	if len(exploitUniqueIDs) == 0 {
		return map[string][]models.Exploit{}, nil
	}

	m := map[string]*redis.StringStringMapCmd{}
	pipe := r.conn.Pipeline()
	for _, exploitUniqueID := range exploitUniqueIDs {
		m[exploitUniqueID] = pipe.HGetAll(ctx, fmt.Sprintf(exploitDBIDKeyFormat, exploitUniqueID))
	}
	if _, err := pipe.Exec(ctx); err != nil {
		return nil, xerrors.Errorf("Failed to exec pipeline. err: %w", err)
	}

	exploits := map[string][]models.Exploit{}
	for exploitUniqueID, cmd := range m {
		results, err := cmd.Result()
		if err != nil {
			return nil, xerrors.Errorf("Failed to get Exploit by ExploitUniqueID. err: %w", err)
		}

		es := []models.Exploit{}
		for _, result := range results {
			var exploit models.Exploit
			if err := json.Unmarshal([]byte(result), &exploit); err != nil {
				return nil, xerrors.Errorf("Failed to Unmarshal json. err: %w", err)
			}
			es = append(es, exploit)
		}
		if len(es) > 0 {
			exploits[exploitUniqueID] = es
		}
	}
	return exploits, nil
}

// GetExploitMultiByCveID :
func (r *RedisDriver) GetExploitMultiByCveID(cveIDs []string) (map[string][]models.Exploit, error) {
	ctx := context.Background()

	if len(cveIDs) == 0 {
		return map[string][]models.Exploit{}, nil
	}

	m := map[string]*redis.StringSliceCmd{}
	pipe := r.conn.Pipeline()
	for _, cveID := range cveIDs {
		m[cveID] = pipe.SMembers(ctx, fmt.Sprintf(cveIDKeyFormat, cveID))
	}
	if _, err := pipe.Exec(ctx); err != nil {
		return nil, xerrors.Errorf("Failed to exec pipeline. err: %w", err)
	}

	mm := map[string][]*redis.StringCmd{}
	pipe = r.conn.Pipeline()
	for cveID, cmd := range m {
		exploitIDs, err := cmd.Result()
		if err != nil {
			return nil, xerrors.Errorf("Failed to get Exploit by CVEID. err: %w", err)
		}

		for _, exploitID := range exploitIDs {
			mm[cveID] = append(mm[cveID], pipe.HGet(ctx, fmt.Sprintf(exploitDBIDKeyFormat, exploitID), cveID))
		}
	}
	if _, err := pipe.Exec(ctx); err != nil {
		return nil, xerrors.Errorf("Failed to exec pipeline. err: %w", err)
	}

	exploits := map[string][]models.Exploit{}
	for cveID, cmds := range mm {
		es := []models.Exploit{}
		for _, cmd := range cmds {
			str, err := cmd.Result()
			if err != nil {
				return nil, xerrors.Errorf("Failed to HGet. err: %w", err)
			}

			var exploit models.Exploit
			if err := json.Unmarshal([]byte(str), &exploit); err != nil {
				return nil, xerrors.Errorf("Failed to Unmarshal json. err: %w", err)
			}
			es = append(es, exploit)
		}
		if len(es) > 0 {
			exploits[cveID] = es
		}
	}
	return exploits, nil
}

// InsertExploit :
func (r *RedisDriver) InsertExploit(exploitType models.ExploitType, exploits []models.Exploit) (err error) {
	ctx := context.Background()
	batchSize := viper.GetInt("batch-size")
	if batchSize < 1 {
		return xerrors.Errorf("Failed to set batch-size. err: batch-size option is not set properly")
	}

	// newDeps, oldDeps: {"ExploitUniqueID": {"CVEID": {}}}
	newDeps := map[string]map[string]struct{}{}
	oldDepsStr, err := r.conn.HGet(ctx, depKey, string(exploitType)).Result()
	if err != nil {
		if !errors.Is(err, redis.Nil) {
			return xerrors.Errorf("Failed to Get key: %s. err: %w", depKey, err)
		}
		oldDepsStr = "{}"
	}
	var oldDeps map[string]map[string]struct{}
	if err := json.Unmarshal([]byte(oldDepsStr), &oldDeps); err != nil {
		return xerrors.Errorf("Failed to unmarshal JSON. err: %w", err)
	}

	bar := pb.StartNew(len(exploits))
	var noCveIDExploitCount, cveIDExploitCount int
	for idx := range chunkSlice(len(exploits), batchSize) {
		pipe := r.conn.Pipeline()
		for _, exploit := range exploits[idx.From:idx.To] {
			j, err := json.Marshal(exploit)
			if err != nil {
				return xerrors.Errorf("Failed to marshal json. err: %w", err)
			}

			_ = pipe.HSet(ctx, fmt.Sprintf(exploitDBIDKeyFormat, exploit.ExploitUniqueID), exploit.CveID, string(j))
			if _, ok := newDeps[exploit.ExploitUniqueID]; !ok {
				newDeps[exploit.ExploitUniqueID] = map[string]struct{}{}
			}

			if exploit.CveID != "" {
				_ = pipe.SAdd(ctx, fmt.Sprintf(cveIDKeyFormat, exploit.CveID), exploit.ExploitUniqueID)
				cveIDExploitCount++
			} else {
				noCveIDExploitCount++
			}

			newDeps[exploit.ExploitUniqueID][exploit.CveID] = struct{}{}
			if _, ok := oldDeps[exploit.ExploitUniqueID]; ok {
				delete(oldDeps[exploit.ExploitUniqueID], exploit.CveID)
				if len(oldDeps[exploit.ExploitUniqueID]) == 0 {
					delete(oldDeps, exploit.ExploitUniqueID)
				}
			}
		}
		if _, err = pipe.Exec(ctx); err != nil {
			return xerrors.Errorf("Failed to exec pipeline. err: %w", err)
		}
		bar.Add(idx.To - idx.From)
	}
	bar.Finish()

	pipe := r.conn.Pipeline()
	for eid, cves := range oldDeps {
		for cveid := range cves {
			if cveid != "" {
				_ = pipe.SRem(ctx, fmt.Sprintf(cveIDKeyFormat, cveid), eid)
			}
			_ = pipe.HDel(ctx, fmt.Sprintf(exploitDBIDKeyFormat, eid), cveid)
		}
	}
	newDepsJSON, err := json.Marshal(newDeps)
	if err != nil {
		return xerrors.Errorf("Failed to Marshal JSON. err: %w", err)
	}
	_ = pipe.HSet(ctx, depKey, string(exploitType), string(newDepsJSON))
	if _, err = pipe.Exec(ctx); err != nil {
		return xerrors.Errorf("Failed to exec pipeline. err: %w", err)
	}

	log15.Info("No CveID Exploit Count", "count", noCveIDExploitCount)
	log15.Info("CveID Exploit Count", "count", cveIDExploitCount)
	return nil
}

// IsExploitModelV1 determines if the DB was created at the time of go-exploitdb Model v1
func (r *RedisDriver) IsExploitModelV1() (bool, error) {
	ctx := context.Background()

	exists, err := r.conn.Exists(ctx, fetchMetaKey).Result()
	if err != nil {
		return false, xerrors.Errorf("Failed to Exists. err: %w", err)
	}
	if exists == 0 {
		keys, _, err := r.conn.Scan(ctx, 0, "EDB#*", 1).Result()
		if err != nil {
			return false, xerrors.Errorf("Failed to Scan. err: %w", err)
		}
		if len(keys) == 0 {
			return false, nil
		}
		return true, nil
	}

	return false, nil
}

// GetFetchMeta get FetchMeta from Database
func (r *RedisDriver) GetFetchMeta() (*models.FetchMeta, error) {
	ctx := context.Background()

	exists, err := r.conn.Exists(ctx, fetchMetaKey).Result()
	if err != nil {
		return nil, xerrors.Errorf("Failed to Exists. err: %w", err)
	}
	if exists == 0 {
		return &models.FetchMeta{ExploitRevision: config.Revision, SchemaVersion: models.LatestSchemaVersion, LastFetchedAt: time.Date(1000, time.January, 1, 0, 0, 0, 0, time.UTC)}, nil
	}

	revision, err := r.conn.HGet(ctx, fetchMetaKey, "Revision").Result()
	if err != nil {
		return nil, xerrors.Errorf("Failed to HGet Revision. err: %w", err)
	}

	verstr, err := r.conn.HGet(ctx, fetchMetaKey, "SchemaVersion").Result()
	if err != nil {
		return nil, xerrors.Errorf("Failed to HGet SchemaVersion. err: %w", err)
	}
	version, err := strconv.ParseUint(verstr, 10, 8)
	if err != nil {
		return nil, xerrors.Errorf("Failed to ParseUint. err: %w", err)
	}

	datestr, err := r.conn.HGet(ctx, fetchMetaKey, "LastFetchedAt").Result()
	if err != nil {
		if !errors.Is(err, redis.Nil) {
			return nil, xerrors.Errorf("Failed to HGet LastFetchedAt. err: %w", err)
		}
		datestr = time.Date(1000, time.January, 1, 0, 0, 0, 0, time.UTC).Format(time.RFC3339)
	}
	date, err := time.Parse(time.RFC3339, datestr)
	if err != nil {
		return nil, xerrors.Errorf("Failed to Parse date. err: %w", err)
	}

	return &models.FetchMeta{ExploitRevision: revision, SchemaVersion: uint(version), LastFetchedAt: date}, nil
}

// UpsertFetchMeta upsert FetchMeta to Database
func (r *RedisDriver) UpsertFetchMeta(fetchMeta *models.FetchMeta) error {
	return r.conn.HSet(context.Background(), fetchMetaKey, map[string]interface{}{"Revision": config.Revision, "SchemaVersion": models.LatestSchemaVersion, "LastFetchedAt": fetchMeta.LastFetchedAt}).Err()
}