b90edf317c
Admin-stored API key with env fallback; scan drive/mail/IMAP uploads. Fail-open if VT down, 422 on malware; migration for virus_scan_status.
110 lines
2.8 KiB
Go
110 lines
2.8 KiB
Go
package virustotal
|
|
|
|
import (
|
|
"context"
|
|
"crypto/sha256"
|
|
"encoding/hex"
|
|
"encoding/json"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strings"
|
|
"testing"
|
|
)
|
|
|
|
func TestScannerLookupMalicious(t *testing.T) {
|
|
sha := sha256.Sum256([]byte("evil"))
|
|
shaHex := hex.EncodeToString(sha[:])
|
|
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
if r.Method == http.MethodGet && strings.Contains(r.URL.Path, "/files/"+shaHex) {
|
|
_ = json.NewEncoder(w).Encode(map[string]any{
|
|
"data": map[string]any{
|
|
"attributes": map[string]any{
|
|
"last_analysis_stats": map[string]any{
|
|
"malicious": 2,
|
|
"suspicious": 0,
|
|
},
|
|
},
|
|
},
|
|
})
|
|
return
|
|
}
|
|
http.NotFound(w, r)
|
|
}))
|
|
defer srv.Close()
|
|
|
|
sc := NewScanner("test-key", nil)
|
|
sc.client.baseURL = srv.URL + "/api/v3"
|
|
|
|
_, err := sc.ScanBytes(context.Background(), "evil.bin", []byte("evil"), shaHex)
|
|
if err == nil {
|
|
t.Fatal("expected ErrMalicious")
|
|
}
|
|
if err != ErrMalicious {
|
|
t.Fatalf("err = %v, want ErrMalicious", err)
|
|
}
|
|
}
|
|
|
|
func TestScannerUploadAndPollClean(t *testing.T) {
|
|
pollCount := 0
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
switch {
|
|
case r.Method == http.MethodGet && strings.HasSuffix(r.URL.Path, "/files/"+SHA256Hex([]byte("clean"))):
|
|
http.NotFound(w, r)
|
|
case r.Method == http.MethodPost && strings.HasSuffix(r.URL.Path, "/files"):
|
|
_ = json.NewEncoder(w).Encode(map[string]any{
|
|
"data": map[string]any{"id": "analysis-1"},
|
|
})
|
|
case r.Method == http.MethodGet && strings.Contains(r.URL.Path, "/analyses/analysis-1"):
|
|
pollCount++
|
|
status := "queued"
|
|
if pollCount >= 2 {
|
|
status = "completed"
|
|
}
|
|
_ = json.NewEncoder(w).Encode(map[string]any{
|
|
"data": map[string]any{
|
|
"attributes": map[string]any{
|
|
"status": status,
|
|
"stats": map[string]any{
|
|
"malicious": 0,
|
|
"suspicious": 0,
|
|
},
|
|
},
|
|
},
|
|
})
|
|
default:
|
|
http.NotFound(w, r)
|
|
}
|
|
}))
|
|
defer srv.Close()
|
|
|
|
sc := NewScanner("test-key", nil)
|
|
sc.client.baseURL = srv.URL + "/api/v3"
|
|
|
|
result, err := sc.ScanBytes(context.Background(), "clean.txt", []byte("clean"), "")
|
|
if err != nil {
|
|
t.Fatalf("ScanBytes: %v", err)
|
|
}
|
|
if result.Status != "clean" {
|
|
t.Fatalf("status = %q, want clean", result.Status)
|
|
}
|
|
}
|
|
|
|
func TestScannerFailOpenOnUnavailable(t *testing.T) {
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
http.Error(w, "down", http.StatusServiceUnavailable)
|
|
}))
|
|
defer srv.Close()
|
|
|
|
sc := NewScanner("test-key", nil)
|
|
sc.client.baseURL = srv.URL + "/api/v3"
|
|
|
|
result, err := sc.ScanBytes(context.Background(), "file.bin", []byte("payload"), "")
|
|
if err != nil {
|
|
t.Fatalf("expected fail-open, got err %v", err)
|
|
}
|
|
if result.Status != "skipped" {
|
|
t.Fatalf("status = %q, want skipped", result.Status)
|
|
}
|
|
}
|