github.com/khulnasoft-lab/tunnel-db@v0.0.0-20231117205118-74e1113bd007/pkg/vulndb/db.go (about)

     1  package vulndb
     2  
     3  import (
     4  	"log"
     5  	"time"
     6  
     7  	bolt "go.etcd.io/bbolt"
     8  	"golang.org/x/xerrors"
     9  	"k8s.io/utils/clock"
    10  
    11  	"github.com/khulnasoft-lab/tunnel-db/pkg/db"
    12  	"github.com/khulnasoft-lab/tunnel-db/pkg/metadata"
    13  	"github.com/khulnasoft-lab/tunnel-db/pkg/types"
    14  	"github.com/khulnasoft-lab/tunnel-db/pkg/vulnsrc"
    15  	"github.com/khulnasoft-lab/tunnel-db/pkg/vulnsrc/vulnerability"
    16  )
    17  
    18  type VulnDB interface {
    19  	Build(targets []string) error
    20  }
    21  
    22  type TunnelDB struct {
    23  	dbc            db.Config
    24  	metadata       metadata.Client
    25  	vulnClient     vulnerability.Vulnerability
    26  	vulnSrcs       map[types.SourceID]vulnsrc.VulnSrc
    27  	cacheDir       string
    28  	updateInterval time.Duration
    29  	clock          clock.Clock
    30  }
    31  
    32  type Option func(*TunnelDB)
    33  
    34  func WithClock(clock clock.Clock) Option {
    35  	return func(core *TunnelDB) {
    36  		core.clock = clock
    37  	}
    38  }
    39  
    40  func WithVulnSrcs(srcs map[types.SourceID]vulnsrc.VulnSrc) Option {
    41  	return func(core *TunnelDB) {
    42  		core.vulnSrcs = srcs
    43  	}
    44  }
    45  
    46  func New(cacheDir string, updateInterval time.Duration, opts ...Option) *TunnelDB {
    47  	// Initialize map
    48  	vulnSrcs := map[types.SourceID]vulnsrc.VulnSrc{}
    49  	for _, v := range vulnsrc.All {
    50  		vulnSrcs[v.Name()] = v
    51  	}
    52  
    53  	dbc := db.Config{}
    54  	tdb := &TunnelDB{
    55  		dbc:            dbc,
    56  		metadata:       metadata.NewClient(cacheDir),
    57  		vulnClient:     vulnerability.New(dbc),
    58  		vulnSrcs:       vulnSrcs,
    59  		cacheDir:       cacheDir,
    60  		updateInterval: updateInterval,
    61  		clock:          clock.RealClock{},
    62  	}
    63  
    64  	for _, opt := range opts {
    65  		opt(tdb)
    66  	}
    67  
    68  	return tdb
    69  }
    70  
    71  func (t TunnelDB) Insert(targets []string) error {
    72  	log.Println("Updating vulnerability database...")
    73  	for _, target := range targets {
    74  		src, ok := t.vulnSrc(target)
    75  		if !ok {
    76  			return xerrors.Errorf("%s is not supported", target)
    77  		}
    78  		log.Printf("Updating %s data...\n", target)
    79  
    80  		if err := src.Update(t.cacheDir); err != nil {
    81  			return xerrors.Errorf("%s update error: %w", target, err)
    82  		}
    83  	}
    84  
    85  	md := metadata.Metadata{
    86  		Version:    db.SchemaVersion,
    87  		NextUpdate: t.clock.Now().UTC().Add(t.updateInterval),
    88  		UpdatedAt:  t.clock.Now().UTC(),
    89  	}
    90  
    91  	if err := t.metadata.Update(md); err != nil {
    92  		return xerrors.Errorf("metadata update error: %w", err)
    93  	}
    94  
    95  	return nil
    96  }
    97  
    98  func (t TunnelDB) Build(targets []string) error {
    99  	// Insert all security advisories
   100  	if err := t.Insert(targets); err != nil {
   101  		return xerrors.Errorf("insert error: %w", err)
   102  	}
   103  
   104  	// Remove unnecessary details
   105  	if err := t.optimize(); err != nil {
   106  		return xerrors.Errorf("optimize error: %w", err)
   107  	}
   108  
   109  	// Remove unnecessary buckets
   110  	if err := t.cleanup(); err != nil {
   111  		return xerrors.Errorf("cleanup error: %w", err)
   112  	}
   113  
   114  	return nil
   115  }
   116  
   117  func (t TunnelDB) vulnSrc(target string) (vulnsrc.VulnSrc, bool) {
   118  	for _, src := range t.vulnSrcs {
   119  		if target == string(src.Name()) {
   120  			return src, true
   121  		}
   122  	}
   123  	return nil, false
   124  }
   125  
   126  func (t TunnelDB) optimize() error {
   127  	// NVD also contains many vulnerabilities that are not related to OS packages or language-specific packages.
   128  	// Tunnel DB will not store them so that it could reduce the database size.
   129  	// This bucket has only vulnerability IDs provided by vendors. They must be stored.
   130  	err := t.dbc.ForEachVulnerabilityID(func(tx *bolt.Tx, cveID string) error {
   131  		details := t.vulnClient.GetDetails(cveID)
   132  		if t.vulnClient.IsRejected(details) {
   133  			return nil
   134  		}
   135  
   136  		if err := t.dbc.SaveAdvisoryDetails(tx, cveID); err != nil {
   137  			return xerrors.Errorf("failed to save advisories: %w", err)
   138  		}
   139  
   140  		if len(details) == 0 {
   141  			return nil
   142  		}
   143  
   144  		vuln := t.vulnClient.Normalize(details)
   145  		if err := t.dbc.PutVulnerability(tx, cveID, vuln); err != nil {
   146  			return xerrors.Errorf("failed to put vulnerability: %w", err)
   147  		}
   148  
   149  		return nil
   150  	})
   151  
   152  	if err != nil {
   153  		return xerrors.Errorf("failed to iterate severity: %w", err)
   154  	}
   155  
   156  	return nil
   157  }
   158  
   159  func (t TunnelDB) cleanup() error {
   160  	if err := t.dbc.DeleteVulnerabilityIDBucket(); err != nil {
   161  		return xerrors.Errorf("failed to delete severity bucket: %w", err)
   162  	}
   163  
   164  	if err := t.dbc.DeleteVulnerabilityDetailBucket(); err != nil {
   165  		return xerrors.Errorf("failed to delete vulnerability detail bucket: %w", err)
   166  	}
   167  
   168  	if err := t.dbc.DeleteAdvisoryDetailBucket(); err != nil {
   169  		return xerrors.Errorf("failed to delete advisory detail bucket: %w", err)
   170  	}
   171  
   172  	return nil
   173  }