81 lines
1.8 KiB
Go
81 lines
1.8 KiB
Go
package capture
|
|
|
|
import (
|
|
"fmt"
|
|
"mime"
|
|
"net/http"
|
|
"os"
|
|
"path/filepath"
|
|
"strings"
|
|
"sync/atomic"
|
|
"time"
|
|
|
|
"git.gulenok.ru/greenhaze/gaslight/internal/domain"
|
|
)
|
|
|
|
type ImageSaver struct {
|
|
outputDir string
|
|
counter uint64
|
|
}
|
|
|
|
func NewImageSaver(outputDir string) *ImageSaver {
|
|
return &ImageSaver{outputDir: outputDir}
|
|
}
|
|
|
|
func (s *ImageSaver) Save(resp *http.Response, body []byte) error {
|
|
n := atomic.AddUint64(&s.counter, 1)
|
|
ts := time.Now().UTC().Format("20060102T150405.000000000Z")
|
|
host := sanitizeName(domain.HostWithoutPort(resp.Request.Host))
|
|
ext := imageExtension(resp.Header.Get("Content-Type"), resp.Request.URL.Path)
|
|
name := fmt.Sprintf("%s_%06d_%s%s", ts, n, host, ext)
|
|
return os.WriteFile(filepath.Join(s.outputDir, name), body, 0o644)
|
|
}
|
|
|
|
func IsImageResponse(contentType, path string) bool {
|
|
if strings.HasPrefix(strings.ToLower(contentType), "image/") {
|
|
return true
|
|
}
|
|
switch strings.ToLower(filepath.Ext(path)) {
|
|
case ".jpg", ".jpeg", ".png", ".gif", ".bmp", ".webp", ".svg", ".ico", ".avif":
|
|
return true
|
|
default:
|
|
return false
|
|
}
|
|
}
|
|
|
|
func imageExtension(contentType, requestPath string) string {
|
|
if ext := strings.ToLower(filepath.Ext(requestPath)); ext != "" && ext != "." {
|
|
return ext
|
|
}
|
|
|
|
base := strings.ToLower(strings.TrimSpace(strings.Split(contentType, ";")[0]))
|
|
if base != "" {
|
|
if exts, err := mime.ExtensionsByType(base); err == nil && len(exts) > 0 {
|
|
return exts[0]
|
|
}
|
|
}
|
|
return ".img"
|
|
}
|
|
|
|
func sanitizeName(s string) string {
|
|
s = strings.ToLower(strings.TrimSpace(s))
|
|
if s == "" {
|
|
return "unknown-host"
|
|
}
|
|
|
|
var b strings.Builder
|
|
for _, r := range s {
|
|
switch {
|
|
case r >= 'a' && r <= 'z':
|
|
b.WriteRune(r)
|
|
case r >= '0' && r <= '9':
|
|
b.WriteRune(r)
|
|
case r == '.', r == '-', r == '_':
|
|
b.WriteRune(r)
|
|
default:
|
|
b.WriteByte('_')
|
|
}
|
|
}
|
|
return b.String()
|
|
}
|