internal/virustotal: add virustotal_api_usage data source

internal/virustotal/client/client.go

@@ -0,0 +1,108 @@
+package client
+
+import (
+	"context"
+	"encoding/json"
+	"net/http"
+	"net/url"
+	"time"
+
+	"github.com/google/go-querystring/query"
+)
+
+var defaultAPIBaseURL = "https://www.virustotal.com/api/v3"
+
+type Client interface {
+	GetUserAPIUsage(ctx context.Context, req *GetUserAPIUsageReq) (*GetUserAPIUsageRes, error)
+	GetGroupAPIUsage(ctx context.Context, req *GetGroupAPIUsageReq) (*GetGroupAPIUsageRes, error)
+}
+
+type client struct {
+	url string
+	key string
+}
+
+func New(key string) Client {
+	return &client{
+		url: defaultAPIBaseURL,
+		key: key,
+	}
+}
+
+func (c *client) auth(r *http.Request) {
+	r.Header.Set("x-apikey", c.key)
+}
+
+func (c *client) GetUserAPIUsage(ctx context.Context, req *GetUserAPIUsageReq) (*GetUserAPIUsageRes, error) {
+	u, err := url.Parse(c.url + "/users/" + req.User + "/api_usage")
+	if err != nil {
+		return nil, err
+	}
+	q, err := query.Values(req)
+	if err != nil {
+		return nil, err
+	}
+	u.RawQuery = q.Encode()
+	r, err := http.NewRequestWithContext(ctx, http.MethodGet, u.String(), nil)
+	if err != nil {
+		return nil, err
+	}
+	c.auth(r)
+	client := http.Client{
+		Timeout: 15 * time.Second,
+	}
+	res, err := client.Do(r)
+	if err != nil {
+		return nil, err
+	}
+	if res.StatusCode != http.StatusOK {
+		var data ErrorRes
+		if err := json.NewDecoder(res.Body).Decode(&data); err != nil {
+			return nil, err
+		}
+		return nil, data.Error
+	}
+	defer res.Body.Close()
+	var data GetUserAPIUsageRes
+	if err := json.NewDecoder(res.Body).Decode(&data); err != nil {
+		return nil, err
+	}
+	return &data, nil
+}
+
+func (c *client) GetGroupAPIUsage(ctx context.Context, req *GetGroupAPIUsageReq) (*GetGroupAPIUsageRes, error) {
+	u, err := url.Parse(c.url + "/groups/" + req.Group + "/api_usage")
+	if err != nil {
+		return nil, err
+	}
+	q, err := query.Values(req)
+	if err != nil {
+		return nil, err
+	}
+	u.RawQuery = q.Encode()
+	r, err := http.NewRequestWithContext(ctx, http.MethodGet, u.String(), nil)
+	if err != nil {
+		return nil, err
+	}
+	c.auth(r)
+	client := http.Client{
+		Timeout: 15 * time.Second,
+	}
+	res, err := client.Do(r)
+	if err != nil {
+		return nil, err
+	}
+	if res.StatusCode != http.StatusOK {
+		var data ErrorRes
+		if err := json.NewDecoder(res.Body).Decode(&data); err != nil {
+			return nil, err
+		}
+		return nil, data.Error
+	}
+	defer res.Body.Close()
+	var data GetGroupAPIUsageRes
+	if err := json.NewDecoder(res.Body).Decode(&data); err != nil {
+		return nil, err
+	}
+	return &data, nil
+}

internal/virustotal/client/client_test.go

@@ -0,0 +1,204 @@
+package client
+
+import (
+	"context"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/suite"
+)
+
+type ClientTestSuite struct {
+	suite.Suite
+	ctx    context.Context
+	cancel context.CancelFunc
+}
+
+func (s *ClientTestSuite) SetupTest() {
+	s.ctx, s.cancel = context.WithCancel(context.Background())
+}
+
+func (s *ClientTestSuite) TearDownTest() {
+	s.cancel()
+}
+
+func TestClientTestSuite(t *testing.T) {
+	suite.Run(t, new(ClientTestSuite))
+}
+
+func (s *ClientTestSuite) mock(fn http.HandlerFunc, tkn string) (*client, *httptest.Server) {
+	srv := httptest.NewServer(fn)
+	cli := &client{
+		url: srv.URL,
+		key: tkn,
+	}
+	return cli, srv
+}
+
+func (s *ClientTestSuite) TestAuth() {
+	client, srv := s.mock(func(w http.ResponseWriter, r *http.Request) {
+		s.Equal("test_token", r.Header.Get("x-apikey"))
+	}, "test_token")
+	defer srv.Close()
+	client.GetUserAPIUsage(s.ctx, &GetUserAPIUsageReq{User: "test_user"})
+}
+
+func (s *ClientTestSuite) TestGetUserAPIUsageWithQuery() {
+	client, srv := s.mock(func(w http.ResponseWriter, r *http.Request) {
+		s.Equal("test_token", r.Header.Get("x-apikey"))
+		s.Equal("GET", r.Method)
+		s.Equal("/users/test_user/api_usage", r.URL.Path)
+		s.Equal("20240101", r.URL.Query().Get("start_date"))
+		s.Equal("20240103", r.URL.Query().Get("end_date"))
+		w.Write([]byte(`{"data": {
+			"daily": {
+				"2024-01-01": {},
+				"2024-01-02": {},
+				"2024-01-03": {}
+			}}}`))
+	}, "test_token")
+	defer srv.Close()
+	start, err := time.Parse("20060102", "20240101")
+	s.Require().NoError(err)
+	end, err := time.Parse("20060102", "20240103")
+	s.Require().NoError(err)
+	res, err := client.GetUserAPIUsage(s.ctx, &GetUserAPIUsageReq{
+		User:      "test_user",
+		StartDate: &Date{start},
+		EndDate:   &Date{end},
+	})
+	s.Require().NoError(err)
+	s.Equal(map[string]any{
+		"daily": map[string]any{
+			"2024-01-01": map[string]any{},
+			"2024-01-02": map[string]any{},
+			"2024-01-03": map[string]any{},
+		},
+	}, res.Data)
+}
+
+func (s *ClientTestSuite) TestGetUserAPIUsage() {
+	client, srv := s.mock(func(w http.ResponseWriter, r *http.Request) {
+		s.Equal("test_token", r.Header.Get("x-apikey"))
+		s.Equal("GET", r.Method)
+		s.Equal("/users/test_user/api_usage", r.URL.Path)
+		w.Write([]byte(`{"data": {
+			"daily": {
+				"2024-01-01": {},
+				"2024-01-02": {},
+				"2024-01-03": {}
+			}}}`))
+	}, "test_token")
+	defer srv.Close()
+	res, err := client.GetUserAPIUsage(s.ctx, &GetUserAPIUsageReq{
+		User: "test_user",
+	})
+	s.Require().NoError(err)
+	s.Equal(map[string]any{
+		"daily": map[string]any{
+			"2024-01-01": map[string]any{},
+			"2024-01-02": map[string]any{},
+			"2024-01-03": map[string]any{},
+		},
+	}, res.Data)
+}
+
+func (s *ClientTestSuite) TestGetGroupAPIUsageWithQuery() {
+	client, srv := s.mock(func(w http.ResponseWriter, r *http.Request) {
+		s.Equal("test_token", r.Header.Get("x-apikey"))
+		s.Equal("GET", r.Method)
+		s.Equal("/groups/test_group/api_usage", r.URL.Path)
+		s.Equal("20240101", r.URL.Query().Get("start_date"))
+		s.Equal("20240103", r.URL.Query().Get("end_date"))
+		w.Write([]byte(`{"data": {
+			"daily": {
+				"2024-01-01": {},
+				"2024-01-02": {},
+				"2024-01-03": {}
+			}}}`))
+	}, "test_token")
+	defer srv.Close()
+	start, err := time.Parse("20060102", "20240101")
+	s.Require().NoError(err)
+	end, err := time.Parse("20060102", "20240103")
+	s.Require().NoError(err)
+	res, err := client.GetGroupAPIUsage(s.ctx, &GetGroupAPIUsageReq{
+		Group:     "test_group",
+		StartDate: &Date{start},
+		EndDate:   &Date{end},
+	})
+	s.Require().NoError(err)
+	s.Equal(map[string]any{
+		"daily": map[string]any{
+			"2024-01-01": map[string]any{},
+			"2024-01-02": map[string]any{},
+			"2024-01-03": map[string]any{},
+		},
+	}, res.Data)
+}
+
+func (s *ClientTestSuite) TestGetGroupAPIUsage() {
+	client, srv := s.mock(func(w http.ResponseWriter, r *http.Request) {
+		s.Equal("test_token", r.Header.Get("x-apikey"))
+		s.Equal("GET", r.Method)
+		s.Equal("/groups/test_group/api_usage", r.URL.Path)
+		w.Write([]byte(`{"data": {
+			"daily": {
+				"2024-01-01": {},
+				"2024-01-02": {},
+				"2024-01-03": {}
+			}}}`))
+	}, "test_token")
+	defer srv.Close()
+	res, err := client.GetGroupAPIUsage(s.ctx, &GetGroupAPIUsageReq{
+		Group: "test_group",
+	})
+	s.Require().NoError(err)
+	s.Equal(map[string]any{
+		"daily": map[string]any{
+			"2024-01-01": map[string]any{},
+			"2024-01-02": map[string]any{},
+			"2024-01-03": map[string]any{},
+		},
+	}, res.Data)
+}
+
+func (s *ClientTestSuite) TestGetUserAPIUsageError() {
+	client, srv := s.mock(func(w http.ResponseWriter, r *http.Request) {
+		s.Equal("test_token", r.Header.Get("x-apikey"))
+		s.Equal("GET", r.Method)
+		s.Equal("/users/test_user/api_usage", r.URL.Path)
+		w.WriteHeader(http.StatusInternalServerError)
+		w.Write([]byte(`{"error": {
+			"code": "test_code",
+			"message": "test_message"
+		}}`))
+	}, "test_token")
+	defer srv.Close()
+	_, err := client.GetUserAPIUsage(s.ctx, &GetUserAPIUsageReq{
+		User: "test_user",
+	})
+	s.Require().Error(err)
+	s.Equal("test_code: test_message", err.Error())
+}
+
+func (s *ClientTestSuite) TestGetGroupAPIUsageError() {
+	client, srv := s.mock(func(w http.ResponseWriter, r *http.Request) {
+		s.Equal("test_token", r.Header.Get("x-apikey"))
+		s.Equal("GET", r.Method)
+		s.Equal("/groups/test_group/api_usage", r.URL.Path)
+		w.WriteHeader(http.StatusInternalServerError)
+		w.Write([]byte(`{"error": {
+			"code": "test_code",
+			"message": "test_message"
+		}}`))
+	}, "test_token")
+	defer srv.Close()
+	_, err := client.GetGroupAPIUsage(s.ctx, &GetGroupAPIUsageReq{
+		Group: "test_group",
+	})
+	s.Require().Error(err)
+	s.Equal("test_code: test_message", err.Error())
+}

internal/virustotal/client/dto.go

@@ -0,0 +1,49 @@
+package client
+
+import (
+	"fmt"
+	"net/url"
+	"time"
+)
+
+type GetUserAPIUsageReq struct {
+	User      string `url:"-"`
+	StartDate *Date  `url:"start_date,omitempty"`
+	EndDate   *Date  `url:"end_date,omitempty"`
+}
+
+type GetGroupAPIUsageReq struct {
+	Group     string `url:"-"`
+	StartDate *Date  `url:"start_date,omitempty"`
+	EndDate   *Date  `url:"end_date,omitempty"`
+}
+
+type Error struct {
+	Code    string `json:"code"`
+	Message string `json:"message"`
+}
+
+func (e Error) Error() string {
+	return fmt.Sprintf("%s: %s", e.Code, e.Message)
+}
+
+type ErrorRes struct {
+	Error Error `json:"error"`
+}
+
+type GetUserAPIUsageRes struct {
+	Data map[string]any `json:"data"`
+}
+
+type GetGroupAPIUsageRes struct {
+	Data map[string]any `json:"data"`
+}
+
+type Date struct {
+	time.Time
+}
+
+func (d Date) EncodeValues(key string, v *url.Values) error {
+	v.Add(key, d.Time.Format("20060102"))
+	return nil
+}

internal/virustotal/cmd/main.go

@@ -0,0 +1,14 @@
+package main
+
+import (
+	"github.com/blackstork-io/fabric/internal/virustotal"
+	pluginapiv1 "github.com/blackstork-io/fabric/plugin/pluginapi/v1"
+)
+
+var version string
+
+func main() {
+	pluginapiv1.Serve(
+		virustotal.Plugin(version, virustotal.DefaultClientLoader),
+	)
+}