ultisuite-backend/internal/virustotal/client.go

package virustotal

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"mime/multipart"
	"net/http"
	"strings"
	"time"
)

const (
	defaultBaseURL    = "https://www.virustotal.com/api/v3"
	directUploadLimit = 32 * 1024 * 1024
	maxUploadLimit    = 650 * 1024 * 1024
)

type Client struct {
	apiKey     string
	baseURL    string
	httpClient *http.Client
}

func NewClient(apiKey string) *Client {
	return &Client{
		apiKey:  apiKey,
		baseURL: defaultBaseURL,
		httpClient: &http.Client{
			Timeout: 120 * time.Second,
		},
	}
}

func (c *Client) apiURL(path string) string {
	return strings.TrimRight(c.baseURL, "/") + path
}

type analysisStats struct {
	Malicious  int `json:"malicious"`
	Suspicious int `json:"suspicious"`
}

type fileReport struct {
	Data struct {
		Attributes struct {
			LastAnalysisStats analysisStats `json:"last_analysis_stats"`
		} `json:"attributes"`
	} `json:"data"`
}

type analysisResponse struct {
	Data struct {
		Attributes struct {
			Status string        `json:"status"`
			Stats  analysisStats `json:"stats"`
		} `json:"attributes"`
	} `json:"data"`
}

type uploadResponse struct {
	Data struct {
		ID string `json:"id"`
	} `json:"data"`
}

type uploadURLResponse struct {
	Data string `json:"data"`
}

func (c *Client) headers() http.Header {
	h := http.Header{}
	h.Set("Accept", "application/json")
	h.Set("x-apikey", c.apiKey)
	return h
}

func (c *Client) lookupFile(ctx context.Context, sha256 string) (analysisStats, bool, error) {
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, c.apiURL("/files/"+sha256), nil)
	if err != nil {
		return analysisStats{}, false, err
	}
	req.Header = c.headers()

	res, err := c.httpClient.Do(req)
	if err != nil {
		return analysisStats{}, false, err
	}
	defer res.Body.Close()

	if res.StatusCode == http.StatusNotFound {
		return analysisStats{}, false, nil
	}
	if res.StatusCode >= 500 || res.StatusCode == http.StatusTooManyRequests {
		return analysisStats{}, false, fmt.Errorf("virustotal lookup unavailable: %d", res.StatusCode)
	}
	if res.StatusCode >= 400 {
		return analysisStats{}, false, nil
	}

	var report fileReport
	if err := json.NewDecoder(res.Body).Decode(&report); err != nil {
		return analysisStats{}, false, err
	}
	return report.Data.Attributes.LastAnalysisStats, true, nil
}

func (c *Client) uploadFile(ctx context.Context, data []byte, filename string) (string, error) {
	if len(data) <= directUploadLimit {
		return c.uploadDirect(ctx, data, filename)
	}
	return c.uploadLarge(ctx, data, filename)
}

func (c *Client) uploadDirect(ctx context.Context, data []byte, filename string) (string, error) {
	var body bytes.Buffer
	w := multipart.NewWriter(&body)
	part, err := w.CreateFormFile("file", filename)
	if err != nil {
		return "", err
	}
	if _, err := part.Write(data); err != nil {
		return "", err
	}
	if err := w.Close(); err != nil {
		return "", err
	}

	req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.apiURL("/files"), &body)
	if err != nil {
		return "", err
	}
	req.Header = c.headers()
	req.Header.Set("Content-Type", w.FormDataContentType())

	res, err := c.httpClient.Do(req)
	if err != nil {
		return "", err
	}
	defer res.Body.Close()

	if res.StatusCode >= 400 {
		b, _ := io.ReadAll(io.LimitReader(res.Body, 4096))
		return "", fmt.Errorf("virustotal upload failed: %d %s", res.StatusCode, string(b))
	}

	var out uploadResponse
	if err := json.NewDecoder(res.Body).Decode(&out); err != nil {
		return "", err
	}
	return out.Data.ID, nil
}

func (c *Client) uploadLarge(ctx context.Context, data []byte, filename string) (string, error) {
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, c.apiURL("/files/upload_url"), nil)
	if err != nil {
		return "", err
	}
	req.Header = c.headers()

	res, err := c.httpClient.Do(req)
	if err != nil {
		return "", err
	}
	defer res.Body.Close()

	if res.StatusCode >= 400 {
		b, _ := io.ReadAll(io.LimitReader(res.Body, 4096))
		return "", fmt.Errorf("virustotal upload_url failed: %d %s", res.StatusCode, string(b))
	}

	var urlResp uploadURLResponse
	if err := json.NewDecoder(res.Body).Decode(&urlResp); err != nil {
		return "", err
	}
	if urlResp.Data == "" {
		return "", fmt.Errorf("virustotal upload_url empty")
	}

	var body bytes.Buffer
	w := multipart.NewWriter(&body)
	part, err := w.CreateFormFile("file", filename)
	if err != nil {
		return "", err
	}
	if _, err := part.Write(data); err != nil {
		return "", err
	}
	if err := w.Close(); err != nil {
		return "", err
	}

	upReq, err := http.NewRequestWithContext(ctx, http.MethodPost, urlResp.Data, &body)
	if err != nil {
		return "", err
	}
	upReq.Header = c.headers()
	upReq.Header.Set("Content-Type", w.FormDataContentType())

	upRes, err := c.httpClient.Do(upReq)
	if err != nil {
		return "", err
	}
	defer upRes.Body.Close()

	if upRes.StatusCode >= 400 {
		b, _ := io.ReadAll(io.LimitReader(upRes.Body, 4096))
		return "", fmt.Errorf("virustotal large upload failed: %d %s", upRes.StatusCode, string(b))
	}

	var out uploadResponse
	if err := json.NewDecoder(upRes.Body).Decode(&out); err != nil {
		return "", err
	}
	return out.Data.ID, nil
}

func (c *Client) pollAnalysis(ctx context.Context, analysisID string, timeout time.Duration) (analysisStats, error) {
	deadline := time.Now().Add(timeout)
	for {
		if ctx.Err() != nil {
			return analysisStats{}, ctx.Err()
		}
		if time.Now().After(deadline) {
			return analysisStats{}, fmt.Errorf("virustotal analysis timeout")
		}

		req, err := http.NewRequestWithContext(ctx, http.MethodGet, c.apiURL("/analyses/"+analysisID), nil)
		if err != nil {
			return analysisStats{}, err
		}
		req.Header = c.headers()

		res, err := c.httpClient.Do(req)
		if err != nil {
			return analysisStats{}, err
		}

		var out analysisResponse
		decodeErr := json.NewDecoder(res.Body).Decode(&out)
		res.Body.Close()
		if decodeErr != nil {
			return analysisStats{}, decodeErr
		}

		if res.StatusCode >= 500 || res.StatusCode == http.StatusTooManyRequests {
			return analysisStats{}, fmt.Errorf("virustotal poll unavailable: %d", res.StatusCode)
		}

		status := out.Data.Attributes.Status
		if status == "completed" {
			return out.Data.Attributes.Stats, nil
		}

		select {
		case <-ctx.Done():
			return analysisStats{}, ctx.Err()
		case <-time.After(2 * time.Second):
		}
	}
}

func statsMalicious(stats analysisStats) bool {
	return stats.Malicious > 0
}