Skip to content

Commit

Permalink
simplify.
Browse files Browse the repository at this point in the history
  • Loading branch information
@jonhadfield
jonhadfield committed Apr 29, 2024
1 parent fcc4c5f commit 5e7d0bc
Show file tree
Hide file tree
Showing 11 changed files with 124 additions and 180 deletions.
150 changes: 38 additions & 112 deletions process/process.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,143 +2,69 @@ package process

import (
"fmt"
"github.com/jonhadfield/ipscout/providers"
"github.com/jonhadfield/ipscout/providers/abuseipdb"
"github.com/jonhadfield/ipscout/providers/annotated"
"github.com/jonhadfield/ipscout/providers/aws"
"github.com/jonhadfield/ipscout/providers/azure"
"github.com/jonhadfield/ipscout/providers/digitalocean"
"github.com/jonhadfield/ipscout/providers/ptr"
"os"
"path/filepath"
"strings"
"sync"
"time"

"github.com/jonhadfield/ipscout/providers/annotated"
"github.com/jonhadfield/ipscout/providers/azure"
"github.com/jonhadfield/ipscout/providers/ptr"

"github.com/briandowns/spinner"
"github.com/jedib0t/go-pretty/v6/table"
"github.com/jonhadfield/ipscout/cache"
"github.com/jonhadfield/ipscout/config"
"github.com/jonhadfield/ipscout/present"
"github.com/jonhadfield/ipscout/providers/abuseipdb"
"github.com/jonhadfield/ipscout/providers/aws"
"github.com/jonhadfield/ipscout/providers/criminalip"
"github.com/jonhadfield/ipscout/providers/digitalocean"
"github.com/jonhadfield/ipscout/providers/ipurl"
"github.com/jonhadfield/ipscout/providers/shodan"
"github.com/mitchellh/go-homedir"
"golang.org/x/sync/errgroup"
)

type ProviderClient interface {
Enabled() bool
GetConfig() *config.Config
Initialise() error
FindHost() ([]byte, error)
CreateTable([]byte) (*table.Writer, error)
type Provider struct {
Name string
Enabled bool
APIKey string
NewClient func(c config.Config) (providers.ProviderClient, error)
}

func getProviderClients(c config.Config) (map[string]ProviderClient, error) {
runners := make(map[string]ProviderClient)

c.Logger.Info("creating provider clients")

if c.Providers.AbuseIPDB.APIKey != "" || c.Providers.AbuseIPDB.Enabled || c.UseTestData {
AbuseIPDBClient, err := abuseipdb.NewProviderClient(c)
if err != nil {
return nil, fmt.Errorf("error creating abuseipdb client: %w", err)
}

if AbuseIPDBClient != nil {
runners[abuseipdb.ProviderName] = AbuseIPDBClient
}
}

if len(c.Providers.Annotated.Paths) > 0 || c.Providers.Annotated.Enabled || c.UseTestData {
AnnotatedClient, err := annotated.NewProviderClient(c)
if err != nil {
return nil, fmt.Errorf("error creating abuseipdb client: %w", err)
}

if AnnotatedClient != nil {
runners[abuseipdb.ProviderName] = AnnotatedClient
}
}

if c.Providers.AWS.Enabled || c.UseTestData {
awsIPClient, err := aws.NewProviderClient(c)
if err != nil {
return nil, fmt.Errorf("error creating aws client: %w", err)
}

if awsIPClient != nil {
runners[aws.ProviderName] = awsIPClient
}
}

if c.Providers.Azure.Enabled || c.UseTestData {
azureIPClient, err := azure.NewProviderClient(c)
if err != nil {
return nil, fmt.Errorf("error creating azure client: %w", err)
}

if azureIPClient != nil {
runners[azure.ProviderName] = azureIPClient
}
}

if c.Providers.CriminalIP.APIKey != "" || c.Providers.CriminalIP.Enabled || c.UseTestData {
criminalIPClient, err := criminalip.NewProviderClient(c)
if err != nil {
return nil, fmt.Errorf("error creating criminalip client: %w", err)
}

if criminalIPClient != nil {
runners[criminalip.ProviderName] = criminalIPClient
}
}

if c.Providers.DigitalOcean.Enabled || c.UseTestData {
digitaloceanIPClient, err := digitalocean.NewProviderClient(c)
if err != nil {
return nil, fmt.Errorf("error creating digitalocean client: %w", err)
}

if digitaloceanIPClient != nil {
runners[digitalocean.ProviderName] = digitaloceanIPClient
}
}

if c.Providers.IPURL.Enabled || c.UseTestData {
IPURLClient, err := ipurl.NewProviderClient(c)
if err != nil {
return nil, fmt.Errorf("error creating ipurl client: %w", err)
}

if IPURLClient != nil {
runners[ipurl.ProviderName] = IPURLClient
}
}

if c.Providers.Shodan.APIKey != "" || c.Providers.Shodan.Enabled || c.UseTestData {
shodanClient, err := shodan.NewProviderClient(c)
if err != nil {
return nil, fmt.Errorf("error creating shodan client: %w", err)
}

runners[shodan.ProviderName] = shodanClient
func getProviderClients(c config.Config) (map[string]providers.ProviderClient, error) {
runners := make(map[string]providers.ProviderClient)

pros := []Provider{
{Name: abuseipdb.ProviderName, Enabled: c.Providers.AbuseIPDB.Enabled, APIKey: c.Providers.AbuseIPDB.APIKey, NewClient: abuseipdb.NewClient},
{Name: annotated.ProviderName, Enabled: c.Providers.Annotated.Enabled, APIKey: "", NewClient: annotated.NewProviderClient},
{Name: aws.ProviderName, Enabled: c.Providers.AWS.Enabled, APIKey: "", NewClient: aws.NewProviderClient},
{Name: azure.ProviderName, Enabled: c.Providers.Azure.Enabled, APIKey: "", NewClient: azure.NewProviderClient},
{Name: criminalip.ProviderName, Enabled: c.Providers.CriminalIP.Enabled, APIKey: c.Providers.CriminalIP.APIKey, NewClient: criminalip.NewProviderClient},
{Name: digitalocean.ProviderName, Enabled: c.Providers.DigitalOcean.Enabled, APIKey: "", NewClient: digitalocean.NewProviderClient},
{Name: ipurl.ProviderName, Enabled: c.Providers.IPURL.Enabled, APIKey: "", NewClient: ipurl.NewProviderClient},
{Name: shodan.ProviderName, Enabled: c.Providers.Shodan.Enabled, APIKey: c.Providers.Shodan.APIKey, NewClient: shodan.NewProviderClient},
{Name: ptr.ProviderName, Enabled: c.Providers.PTR.Enabled, APIKey: "", NewClient: ptr.NewProviderClient},
}

if c.Providers.PTR.Enabled {
ptrClient, err := ptr.NewProviderClient(c)
if err != nil {
return nil, fmt.Errorf("error creating ptr client: %w", err)
for _, provider := range pros {
if provider.Enabled || c.UseTestData || provider.APIKey != "" {
client, err := provider.NewClient(c)
if err != nil {
return nil, fmt.Errorf("error creating %s client: %w", provider.Name, err)
}
if client != nil {
runners[provider.Name] = client
}
}

runners[ptr.ProviderName] = ptrClient
}

return runners, nil
}

func getEnabledProviders(runners map[string]ProviderClient) []string {
func getEnabledProviders(runners map[string]providers.ProviderClient) []string {
keys := make([]string, len(runners))

i := 0
Expand Down Expand Up @@ -258,7 +184,7 @@ func (p *Processor) Run() {
}
}

func initialiseProviders(runners map[string]ProviderClient, hideProgress bool) error {
func initialiseProviders(runners map[string]providers.ProviderClient, hideProgress bool) error {
var err error

var g errgroup.Group
Expand Down Expand Up @@ -317,7 +243,7 @@ type generateTablesResults struct {
m []*table.Writer
}

func findHosts(runners map[string]ProviderClient, hideProgress bool) (*findHostsResults, error) {
func findHosts(runners map[string]providers.ProviderClient, hideProgress bool) (*findHostsResults, error) {
var results findHostsResults

results.Lock()
Expand Down Expand Up @@ -363,7 +289,7 @@ func findHosts(runners map[string]ProviderClient, hideProgress bool) (*findHosts
return &results, nil
}

func generateTables(conf *config.Config, runners map[string]ProviderClient, results *findHostsResults) ([]*table.Writer, error) {
func generateTables(conf *config.Config, runners map[string]providers.ProviderClient, results *findHostsResults) ([]*table.Writer, error) {
var tables generateTablesResults

var w sync.WaitGroup
Expand Down
32 changes: 16 additions & 16 deletions providers/abuseipdb/abuseipdb.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -31,10 +31,10 @@ const (
ResultTTL = time.Duration(12 * time.Hour)
)

type Client struct {
Config Config
HTTPClient *retryablehttp.Client
}
//type Client struct {
// Config Config
// HTTPClient *retryablehttp.Client
//}

type Config struct {
_ struct{}
Expand All @@ -43,19 +43,19 @@ type Config struct {
APIKey string
}

func NewProviderClient(c config.Config) (*ProviderClient, error) {
func NewClient(c config.Config) (providers.ProviderClient, error) {
c.Logger.Debug("creating abuseipdb client")

tc := &ProviderClient{
tc := Client{
c,
}

return tc, nil
return &tc, nil
}

func (c *Client) GetConfig() *config.Config {
return &c.Config.Config
}
//func (c *Client) GetConfig() *config.Config {
// return &c.Config.Config
//}

func (c *Client) GetData() (result *HostSearchResult, err error) {
result, err = loadResultsFile("abuseipdb/testdata/abuseipdb_google_dns_resp.json")
Expand All @@ -71,19 +71,19 @@ type Provider interface {
CreateTable([]byte) (*table.Writer, error)
}

func (c *ProviderClient) Enabled() bool {
func (c *Client) Enabled() bool {
return c.Config.Providers.AbuseIPDB.Enabled
}

func (c *ProviderClient) GetConfig() *config.Config {
func (c *Client) GetConfig() *config.Config {
return &c.Config
}

type ProviderClient struct {
type Client struct {
config.Config
}

func (c *ProviderClient) Initialise() error {
func (c *Client) Initialise() error {
if c.Cache == nil {
return errors.New("cache not set")
}
Expand All @@ -104,7 +104,7 @@ func (c *ProviderClient) Initialise() error {
return nil
}

func (c *ProviderClient) FindHost() ([]byte, error) {
func (c *Client) FindHost() ([]byte, error) {
start := time.Now()
defer func() {
c.Stats.Mu.Lock()
Expand All @@ -122,7 +122,7 @@ func (c *ProviderClient) FindHost() ([]byte, error) {
return result.Raw, nil
}

func (c *ProviderClient) CreateTable(data []byte) (*table.Writer, error) {
func (c *Client) CreateTable(data []byte) (*table.Writer, error) {
start := time.Now()
defer func() {
c.Stats.Mu.Lock()
Expand Down
2 changes: 1 addition & 1 deletion providers/annotated/annotated.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -50,7 +50,7 @@ func (c *ProviderClient) GetConfig() *config.Config {
return &c.Config
}

func NewProviderClient(c config.Config) (*ProviderClient, error) {
func NewProviderClient(c config.Config) (providers.ProviderClient, error) {
c.Logger.Debug("creating annotated client")

if c.Logger == nil {
Expand Down
2 changes: 1 addition & 1 deletion providers/aws/aws.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -33,7 +33,7 @@ type ProviderClient struct {
config.Config
}

func NewProviderClient(c config.Config) (*ProviderClient, error) {
func NewProviderClient(c config.Config) (providers.ProviderClient, error) {
c.Logger.Debug("creating aws client")

tc := &ProviderClient{
Expand Down
2 changes: 1 addition & 1 deletion providers/azure/azure.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -42,7 +42,7 @@ type ProviderClient struct {
config.Config
}

func NewProviderClient(c config.Config) (*ProviderClient, error) {
func NewProviderClient(c config.Config) (providers.ProviderClient, error) {
c.Logger.Debug("creating azure client")

tc := &ProviderClient{
Expand Down
20 changes: 10 additions & 10 deletions providers/criminalip/criminalip.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -30,25 +30,25 @@ const (
ResultTTL = time.Duration(24 * time.Hour)
)

type ProviderClient struct {
type Client struct {
config.Config
}

func NewProviderClient(c config.Config) (*ProviderClient, error) {
func NewProviderClient(c config.Config) (providers.ProviderClient, error) {
c.Logger.Debug("creating criminalip client")

tc := &ProviderClient{
tc := &Client{
Config: c,
}

return tc, nil
}

func (c *ProviderClient) GetConfig() *config.Config {
func (c *Client) GetConfig() *config.Config {
return &c.Config
}

func (c *ProviderClient) Enabled() bool {
func (c *Client) Enabled() bool {
return c.Config.Providers.CriminalIP.Enabled
}

Expand Down Expand Up @@ -140,7 +140,7 @@ func unmarshalResponse(rBody []byte) (*HostSearchResult, error) {
return res, nil
}

func loadTestData(c *ProviderClient) ([]byte, error) {
func loadTestData(c *Client) ([]byte, error) {
tdf, err := loadResultsFile("providers/criminalip/testdata/criminalip_9_9_9_9_report.json")
if err != nil {
return nil, err
Expand Down Expand Up @@ -229,7 +229,7 @@ func getDomains(domain HostSearchResultDomain) []string {
return domains
}

func (c *ProviderClient) Initialise() error {
func (c *Client) Initialise() error {
if c.Cache == nil {
return errors.New("cache not set")
}
Expand All @@ -246,7 +246,7 @@ func (c *ProviderClient) Initialise() error {
return nil
}

func (c *ProviderClient) FindHost() ([]byte, error) {
func (c *Client) FindHost() ([]byte, error) {
start := time.Now()
defer func() {
c.Stats.Mu.Lock()
Expand Down Expand Up @@ -280,7 +280,7 @@ type GeneratePortDataForTableOutput struct {
matches int
}

func (c *ProviderClient) GenPortDataForTable(in []PortDataEntry) (GeneratePortDataForTableOutput, error) {
func (c *Client) GenPortDataForTable(in []PortDataEntry) (GeneratePortDataForTableOutput, error) {
var err error

var out GeneratePortDataForTableOutput
Expand Down Expand Up @@ -319,7 +319,7 @@ func (c *ProviderClient) GenPortDataForTable(in []PortDataEntry) (GeneratePortDa
return out, nil
}

func (c *ProviderClient) CreateTable(data []byte) (*table.Writer, error) {
func (c *Client) CreateTable(data []byte) (*table.Writer, error) {
start := time.Now()
defer func() {
c.Stats.Mu.Lock()
Expand Down
Loading

0 comments on commit 5e7d0bc

Please sign in to comment.