Skip to content

Commit

Permalink
feat(redis): add expire option (#55)
Browse files Browse the repository at this point in the history
* feat(redis): add expire option

* feat(redis): make the Key persistent when expire is 0

* chore: do not repeat the same string concat
  • Loading branch information
MaineK00n authored Aug 12, 2021
1 parent eb8625d commit 36b7f56
Show file tree
Hide file tree
Showing 3 changed files with 37 additions and 7 deletions.
8 changes: 5 additions & 3 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -44,7 +44,8 @@ $ make install

### Usage: Fetch and Insert Exploit
```bash
$ Fetch the data of exploit
$ go-exploitdb fetch --help
Fetch the data of exploit

Usage:
go-exploitdb fetch [command]
Expand All @@ -55,12 +56,13 @@ Available Commands:
githubrepos Fetch the data of github repos

Flags:
-h, --help help for fetch
--expire uint timeout to set for Redis keys in seconds. If set to 0, the key is persistent.
-h, --help help for fetch

Global Flags:
--config string config file (default is $HOME/.go-exploitdb.yaml)
--dbpath string /path/to/sqlite3 or SQL connection string
--dbtype string Database type to store data in (sqlite3, mysql, postgres, or redis supported)
--dbtype string Database type to store data in (sqlite3, mysql, postgres or redis supported)
--debug debug mode (default: false)
--debug-sql SQL debug mode
--http-proxy string http://proxy-url:port (default: empty)
Expand Down
4 changes: 4 additions & 0 deletions commands/fetch.go
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ package commands

import (
"github.com/spf13/cobra"
"github.com/spf13/viper"
)

// fetchCmd represents the fetch command
Expand All @@ -13,4 +14,7 @@ var fetchCmd = &cobra.Command{

func init() {
RootCmd.AddCommand(fetchCmd)

fetchCmd.PersistentFlags().Uint("expire", 0, "timeout to set for Redis keys in seconds. If set to 0, the key is persistent.")
_ = viper.BindPFlag("expire", fetchCmd.PersistentFlags().Lookup("expire"))
}
32 changes: 28 additions & 4 deletions db/redis.go
Original file line number Diff line number Diff line change
Expand Up @@ -4,10 +4,12 @@ import (
"context"
"encoding/json"
"fmt"
"time"

"github.com/cheggaaa/pb/v3"
"github.com/go-redis/redis/v8"
"github.com/inconshreveable/log15"
"github.com/spf13/viper"
"github.com/vulsio/go-exploitdb/models"
"golang.org/x/xerrors"
)
Expand Down Expand Up @@ -163,13 +165,14 @@ func (r *RedisDriver) GetExploitMultiByCveID(cveIDs []string) (exploitsMap map[s

//InsertExploit :
func (r *RedisDriver) InsertExploit(_ models.ExploitType, exploits []*models.Exploit) (err error) {
expire := viper.GetUint("expire")

ctx := context.Background()
bar := pb.StartNew(len(exploits))

var noCveIDExploitCount, cveIDExploitCount int
for _, exploit := range exploits {
var pipe redis.Pipeliner
pipe = r.conn.Pipeline()
pipe := r.conn.Pipeline()
bar.Increment()

j, err := json.Marshal(exploit)
Expand All @@ -178,9 +181,19 @@ func (r *RedisDriver) InsertExploit(_ models.ExploitType, exploits []*models.Exp
}

if 0 < len(exploit.CveID) {
if result := pipe.HSet(ctx, cveIDPrefix+exploit.CveID, exploit.ExploitUniqueID, string(j)); result.Err() != nil {
key := cveIDPrefix + exploit.CveID
if result := pipe.HSet(ctx, key, exploit.ExploitUniqueID, string(j)); result.Err() != nil {
return fmt.Errorf("Failed to HSet CVE. err: %s", result.Err())
}
if expire > 0 {
if err := pipe.Expire(ctx, key, time.Duration(expire*uint(time.Second))).Err(); err != nil {
return fmt.Errorf("Failed to set Expire to Key. err: %s", err)
}
} else {
if err := pipe.Persist(ctx, key).Err(); err != nil {
return fmt.Errorf("Failed to remove the existing timeout on Key. err: %s", err)
}
}
cveIDExploitCount++
} else {
noCveIDExploitCount++
Expand All @@ -190,9 +203,20 @@ func (r *RedisDriver) InsertExploit(_ models.ExploitType, exploits []*models.Exp
if len(exploit.CveID) == 0 {
exploit.CveID = "NONE"
}
if result := pipe.HSet(ctx, exploitDBIDPrefix+exploit.ExploitUniqueID, exploit.CveID, string(j)); result.Err() != nil {

key := exploitDBIDPrefix + exploit.ExploitUniqueID
if result := pipe.HSet(ctx, key, exploit.CveID, string(j)); result.Err() != nil {
return fmt.Errorf("Failed to HSet Exploit. err: %s", result.Err())
}
if expire > 0 {
if err := pipe.Expire(ctx, key, time.Duration(expire*uint(time.Second))).Err(); err != nil {
return fmt.Errorf("Failed to set Expire to Key. err: %s", err)
}
} else {
if err := pipe.Persist(ctx, key).Err(); err != nil {
return fmt.Errorf("Failed to remove the existing timeout on Key. err: %s", err)
}
}

if _, err = pipe.Exec(ctx); err != nil {
return fmt.Errorf("Failed to exec pipeline. err: %s", err)
Expand Down

0 comments on commit 36b7f56

Please sign in to comment.