hako/internal/archival/archiver/obelisk.go

package archiver

import (
	"bytes"
	"context"
	"crypto/sha256"
	"encoding/hex"
	"fmt"
	"io"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"git.nakama.town/fmartingr/hako/internal/model"
	"git.nakama.town/fmartingr/hako/internal/storage"
	"github.com/go-shiori/obelisk"
)

// ObeliskArchiver implements web page archival using obelisk
type ObeliskArchiver struct {
	archiver *obelisk.Archiver
	mu       sync.RWMutex // Protects archiver state during concurrent access
}

// NewObeliskExtractor creates a new ObeliskArchiver
func NewObeliskExtractor() *ObeliskArchiver {
	archiver := &ObeliskArchiver{
		archiver: &obelisk.Archiver{
			EnableLog:        false,
			EnableVerboseLog: false,
		},
	}

	// Apply default config
	defaultConfig := archiver.GetDefaultConfig().(map[string]any)
	if err := archiver.applyConfigToArchiver(defaultConfig); err != nil {
		// If applying default config fails, use hardcoded fallback values
		archiver.archiver.UserAgent = "Hako/1.0"
		archiver.archiver.RequestTimeout = 5 * time.Minute
		archiver.archiver.MaxRetries = 3
		archiver.archiver.MaxConcurrentDownload = 10
		archiver.archiver.SkipResourceURLError = false
		archiver.archiver.Validate()
	}

	return archiver
}

// Key returns the unique identifier for this archiver
func (e *ObeliskArchiver) Key() string {
	return "obelisk"
}

// Name returns the human-readable name for this archiver
func (e *ObeliskArchiver) Name() string {
	return "Obelisk"
}

// IsEnabled checks if this archiver is available
func (e *ObeliskArchiver) IsEnabled(ctx context.Context) bool {
	// Obelisk is always enabled
	return true
}

// Init initializes the archiver
func (e *ObeliskArchiver) Init() error {
	// Archiver is already initialized in NewObeliskExtractor
	return nil
}

// GetDefaultConfig returns the default configuration
func (e *ObeliskArchiver) GetDefaultConfig() any {
	return map[string]any{
		"timeout":                 "5m",
		"max_retries":             3,
		"max_concurrent_download": 10,
		"user_agent":              "Hako/1.0",
		"skip_resource_url_error": false,
	}
}

// applyConfigToArchiver applies configuration values to the archiver
func (e *ObeliskArchiver) applyConfigToArchiver(config map[string]any) error {
	e.mu.Lock()
	defer e.mu.Unlock()

	// Parse timeout duration
	if timeoutStr, ok := config["timeout"].(string); ok {
		timeout, err := time.ParseDuration(timeoutStr)
		if err != nil {
			return fmt.Errorf("invalid timeout format: %w", err)
		}
		e.archiver.RequestTimeout = timeout
	}

	// Parse max_retries
	if maxRetries, ok := config["max_retries"].(float64); ok {
		e.archiver.MaxRetries = int(maxRetries)
	} else if maxRetries, ok := config["max_retries"].(int); ok {
		e.archiver.MaxRetries = maxRetries
	}

	// Parse max_concurrent_download
	if maxConcurrent, ok := config["max_concurrent_download"].(float64); ok {
		e.archiver.MaxConcurrentDownload = int64(maxConcurrent)
	} else if maxConcurrent, ok := config["max_concurrent_download"].(int64); ok {
		e.archiver.MaxConcurrentDownload = maxConcurrent
	} else if maxConcurrent, ok := config["max_concurrent_download"].(int); ok {
		e.archiver.MaxConcurrentDownload = int64(maxConcurrent)
	}

	// Parse user_agent
	if userAgent, ok := config["user_agent"].(string); ok {
		e.archiver.UserAgent = userAgent
	}

	// Parse skip_resource_url_error
	if skipError, ok := config["skip_resource_url_error"].(bool); ok {
		e.archiver.SkipResourceURLError = skipError
	}

	// Validate archiver after applying config
	e.archiver.Validate()

	return nil
}

// ApplyConfig applies the provided configuration to the archiver
func (e *ObeliskArchiver) ApplyConfig(config map[string]any) error {
	return e.applyConfigToArchiver(config)
}

// Archive performs the obelisk archival
// Note: ApplyConfig should be called before Archive() if using custom configuration
func (e *ObeliskArchiver) Archive(ctx context.Context, link *model.Link, stor storage.Storage) (*ArchiveResult, error) {
	e.mu.RLock()
	archiver := e.archiver
	e.mu.RUnlock()

	// Create obelisk request
	req := obelisk.Request{
		URL: link.URL,
	}

	// Perform the archival
	htmlBytes, title, err := archiver.Archive(ctx, req)
	if err != nil {
		return nil, fmt.Errorf("failed to archive with obelisk: %w", err)
	}

	// Extract page title from HTML
	pageTitle := extractPageTitle(htmlBytes)

	// Use page title from HTML if obelisk's title is missing or looks like a content-type
	if pageTitle != "" && (title == "" || looksLikeContentType(title)) {
		title = pageTitle
	}

	// Generate filename based on title or URL
	filename := "index.html"
	if title != "" && !looksLikeContentType(title) {
		// Sanitize title for filename
		filename = sanitizeFilename(title) + ".html"
	} else {
		// Try to extract filename from URL
		urlPath := filepath.Base(link.URL)
		if urlPath != "." && urlPath != "/" && urlPath != "" {
			// Remove query parameters
			if idx := strings.Index(urlPath, "?"); idx != -1 {
				urlPath = urlPath[:idx]
			}
			if urlPath != "" {
				// Ensure it has .html extension
				if filepath.Ext(urlPath) == "" {
					filename = urlPath + ".html"
				} else {
					filename = urlPath
				}
			}
		}
	}

	// Create a pipe to calculate hash while saving
	pipeReader, pipeWriter := io.Pipe()
	hash := sha256.New()
	multiWriter := io.MultiWriter(pipeWriter, hash)

	// Channel for errors from the goroutine
	errChan := make(chan error, 1)
	var path string
	var fileSize int64

	// Copy to storage in goroutine
	go func() {
		defer func() { _ = pipeWriter.Close() }()
		reader := bytes.NewReader(htmlBytes)
		if _, err := io.Copy(multiWriter, reader); err != nil {
			errChan <- fmt.Errorf("failed to copy data: %w", err)
			return
		}
		errChan <- nil
	}()

	// Save to storage
	path, err = stor.Save(ctx, link.ID, filename, pipeReader)
	if err != nil {
		return nil, fmt.Errorf("failed to save file: %w", err)
	}

	// Check for copy errors
	if copyErr := <-errChan; copyErr != nil {
		return nil, copyErr
	}

	// Get file size
	fileSize = int64(len(htmlBytes))

	// Calculate hash
	hashStr := hex.EncodeToString(hash.Sum(nil))

	// Use page title if available, otherwise use URL
	resultTitle := title
	if resultTitle == "" || looksLikeContentType(resultTitle) {
		if pageTitle != "" {
			resultTitle = pageTitle
		} else {
			resultTitle = link.URL
		}
	}

	result := &ArchiveResult{
		Title: resultTitle,
		Files: []FileInfo{
			{
				Filename:   filename,
				MimeType:   "text/html",
				FileSize:   fileSize,
				HashSha256: hashStr,
				Path:       path,
			},
		},
	}

	return result, nil
}

// extractPageTitle extracts the page title from HTML content
func extractPageTitle(htmlBytes []byte) string {
	htmlStr := string(htmlBytes)

	// Look for <title> tag (case-insensitive)
	titleStart := strings.Index(strings.ToLower(htmlStr), "<title>")
	if titleStart == -1 {
		return ""
	}

	titleStart += len("<title>")
	titleEnd := strings.Index(strings.ToLower(htmlStr[titleStart:]), "</title>")
	if titleEnd == -1 {
		return ""
	}

	title := htmlStr[titleStart : titleStart+titleEnd]
	// Decode HTML entities if needed (basic handling)
	title = strings.ReplaceAll(title, "&nbsp;", " ")
	title = strings.ReplaceAll(title, "&amp;", "&")
	title = strings.ReplaceAll(title, "&lt;", "<")
	title = strings.ReplaceAll(title, "&gt;", ">")
	title = strings.ReplaceAll(title, "&quot;", "\"")
	title = strings.ReplaceAll(title, "&#39;", "'")

	return strings.TrimSpace(title)
}

// looksLikeContentType checks if a string looks like a content-type header value
func looksLikeContentType(s string) bool {
	s = strings.ToLower(strings.TrimSpace(s))
	// Check for common content-type patterns
	return strings.Contains(s, "text/") ||
		strings.Contains(s, "application/") ||
		strings.Contains(s, "charset=") ||
		strings.Contains(s, "content-type")
}

// sanitizeFilename sanitizes a string to be used as a filename
func sanitizeFilename(s string) string {
	// Remove or replace invalid filename characters
	invalid := []string{"/", "\\", ":", "*", "?", "\"", "<", ">", "|", "\n", "\r", "\t"}
	result := s
	for _, char := range invalid {
		result = strings.ReplaceAll(result, char, "_")
	}
	// Limit length
	if len(result) > 200 {
		result = result[:200]
	}
	// Trim whitespace
	result = strings.TrimSpace(result)
	if result == "" {
		result = "index"
	}
	return result
}