system_hosts_manager.go

package main

import (
	"fmt"
	"io"
	"os"
	"strings"
)

const (
	hostsFilePath       = "/etc/hosts"           // The actual system file
	hostsFileBackupPath = "/etc/hosts.dsfree.bak" // Our backup location
	redirectIP          = "127.0.0.1"            // The IP to redirect to
)

// SystemHostsManager is the concrete implementation of HostsManager.
// It's the concrete implementation that directly interacts with the file system.
type SystemHostsManager struct{}

// Apply modifies the real /etc/hosts file.
// Apply is now idempotent. It ensures the hosts file is in the blocked state.
func (s *SystemHostsManager) Apply(domains []string) error {
	// Ensure a backup exists. If not, create one.
	// This makes the operation safe to re-run.
	if _, err := os.Stat(hostsFileBackupPath); os.IsNotExist(err) {
		if err := s.backupHostsFile(); err != nil {
			return fmt.Errorf("failed to create initial backup: %w", err)
		}
	}

	// ALWAYS read from the clean backup, not the potentially tampered file.
	originalContent, err := os.ReadFile(hostsFileBackupPath)
	if err != nil {
		return fmt.Errorf("failed to read backup hosts file: %w", err)
	}

	// Build the new lines to add.
	var newLines []string
	newLines = append(newLines, "\n# --- DSFreeService BLOCKLIST START ---")
	for _, domain := range domains {
		newLines = append(newLines, fmt.Sprintf("%s %s", redirectIP, domain))
	}
	newLines = append(newLines, "# --- DSFreeService BLOCKLIST END ---")

	// Append the new lines to the original content and overwrite the live file.
	// This reverts any manual changes and enforces our state.
	finalContent := string(originalContent) + strings.Join(newLines, "\n")
	if err := os.WriteFile(hostsFilePath, []byte(finalContent), 0644); err != nil {
		return fmt.Errorf("failed to write to hosts file: %w", err)
	}

	return nil
}

// Restore reverts the hosts file from the backup.
func (s *SystemHostsManager) Restore() error {
	// Simply copy the backup file back to the original location.
	source, err := os.Open(hostsFileBackupPath)
	if err != nil {
		// If backup doesn't exist, maybe it was never created. That's ok.
		if os.IsNotExist(err) {
			return nil
		}
		return fmt.Errorf("failed to open backup file: %w", err)
	}
	defer source.Close()

	destination, err := os.Create(hostsFilePath)
	if err != nil {
		return fmt.Errorf("failed to open hosts file for writing: %w", err)
	}
	defer destination.Close()

	_, err = io.Copy(destination, source)
	if err != nil {
		return fmt.Errorf("failed to copy backup to hosts file: %w", err)
	}

	// Clean up the backup file.
	return os.Remove(hostsFileBackupPath)
}

// backupHostsFile is a private helper method.
func (s *SystemHostsManager) backupHostsFile() error {
	// Check if a backup already exists. If so, don't overwrite it.
	// This implies a previous run did not clean up properly.
	if _, err := os.Stat(hostsFileBackupPath); err == nil {
		return fmt.Errorf("backup file %s already exists, please remove it manually", hostsFileBackupPath)
	}

	source, err := os.Open(hostsFilePath)
	if err != nil {
		return err
	}
	defer source.Close()

	destination, err := os.Create(hostsFileBackupPath)
	if err != nil {
		return err
	}
	defer destination.Close()

	_, err = io.Copy(destination, source)
	return err
}