fix: share key retriever across all browsers to avoid repeated prompts (#560)

* fix: share key retriever across all browsers to avoid repeated password prompts
2026-05-19 18:58:03 +02:00 · 2026-04-06 21:57:52 +08:00
parent ccc8643d86
commit a0b4412bf2
8 changed files with 355 additions and 80 deletions
@@ -8,6 +8,7 @@ import (

 	"github.com/moond4rk/hackbrowserdata/browser/chromium"
 	"github.com/moond4rk/hackbrowserdata/browser/firefox"
+	"github.com/moond4rk/hackbrowserdata/crypto/keyretriever"
 	"github.com/moond4rk/hackbrowserdata/log"
 	"github.com/moond4rk/hackbrowserdata/types"
 )
@@ -34,19 +35,27 @@ func PickBrowsers(opts PickOptions) ([]Browser, error) {
 	return pickFromConfigs(platformBrowsers(), opts)
 }

-// pickFromConfigs is the testable core of PickBrowsers.
+// pickFromConfigs is the testable core of PickBrowsers. It iterates over
+// platform browser configs, discovers installed profiles, and injects a
+// shared key retriever into Chromium browsers for decryption.
 func pickFromConfigs(configs []types.BrowserConfig, opts PickOptions) ([]Browser, error) {
 	name := strings.ToLower(opts.Name)
 	if name == "" {
 		name = "all"
 	}

+	// Create a single key retriever shared across all Chromium browsers.
+	// On macOS this avoids repeated password prompts; on other platforms
+	// it's harmless (DPAPI reads Local State per-profile, D-Bus is stateless).
+	retriever := keyretriever.DefaultRetriever(opts.KeychainPassword)
+
 	var browsers []Browser
 	for _, cfg := range configs {
 		if name != "all" && cfg.Key != name {
 			continue
 		}

+		// Override profile directory when targeting a specific browser.
 		if opts.ProfilePath != "" && name != "all" {
 			if cfg.Kind == types.Firefox {
 				cfg.UserDataDir = filepath.Dir(filepath.Clean(opts.ProfilePath))
@@ -55,48 +64,60 @@ func pickFromConfigs(configs []types.BrowserConfig, opts PickOptions) ([]Browser
 			}
 		}

-		if opts.KeychainPassword != "" {
-			cfg.KeychainPassword = opts.KeychainPassword
-		}
-
-		bs, err := newBrowsers(cfg)
+		found, err := newBrowsers(cfg)
 		if err != nil {
 			log.Errorf("browser %s: %v", cfg.Name, err)
 			continue
 		}
-		if len(bs) == 0 {
+		if len(found) == 0 {
 			log.Debugf("browser %s not found at %s", cfg.Name, cfg.UserDataDir)
 			continue
 		}
-		browsers = append(browsers, bs...)
+
+		// Inject the shared key retriever into browsers that need it.
+		// Chromium browsers implement retrieverSetter; Firefox does not.
+		for _, b := range found {
+			if setter, ok := b.(retrieverSetter); ok {
+				setter.SetRetriever(retriever)
+			}
+		}
+		browsers = append(browsers, found...)
 	}
 	return browsers, nil
 }

-// newBrowsers dispatches to the correct engine based on BrowserKind.
+// retrieverSetter is implemented by browsers that need an external key retriever.
+// This allows pickFromConfigs to inject the shared retriever after construction
+// without coupling the Browser interface to Chromium-specific concerns.
+type retrieverSetter interface {
+	SetRetriever(keyretriever.KeyRetriever)
+}
+
+// newBrowsers dispatches to the correct engine based on BrowserKind
+// and converts engine-specific types to the Browser interface.
 func newBrowsers(cfg types.BrowserConfig) ([]Browser, error) {
 	switch cfg.Kind {
 	case types.Chromium, types.ChromiumYandex, types.ChromiumOpera:
-		bs, err := chromium.NewBrowsers(cfg)
+		found, err := chromium.NewBrowsers(cfg)
 		if err != nil {
 			return nil, err
 		}
-		browsers := make([]Browser, len(bs))
-		for i, b := range bs {
-			browsers[i] = b
+		result := make([]Browser, len(found))
+		for i, b := range found {
+			result[i] = b
 		}
-		return browsers, nil
+		return result, nil

 	case types.Firefox:
-		bs, err := firefox.NewBrowsers(cfg)
+		found, err := firefox.NewBrowsers(cfg)
 		if err != nil {
 			return nil, err
 		}
-		browsers := make([]Browser, len(bs))
-		for i, b := range bs {
-			browsers[i] = b
+		result := make([]Browser, len(found))
+		for i, b := range found {
+			result[i] = b
 		}
-		return browsers, nil
+		return result, nil

 	default:
 		return nil, fmt.Errorf("unknown browser kind: %d", cfg.Kind)
@@ -202,6 +202,75 @@ func TestPickFromConfigs_ProfilePath(t *testing.T) {
 	}
 }

+// ---------------------------------------------------------------------------
+// newBrowsers dispatcher
+// ---------------------------------------------------------------------------
+
+func TestNewBrowsersDispatch(t *testing.T) {
+	chromiumDir := t.TempDir()
+	mkFile(t, chromiumDir, "Default", "Preferences")
+	mkFile(t, chromiumDir, "Default", "History")
+
+	firefoxDir := t.TempDir()
+	mkFile(t, firefoxDir, "abc.default", "places.sqlite")
+
+	emptyDir := t.TempDir()
+
+	tests := []struct {
+		name        string
+		cfg         types.BrowserConfig
+		wantLen     int
+		wantName    string
+		wantProfile string
+		wantErr     string
+	}{
+		{
+			name:        "chromium dispatch",
+			cfg:         types.BrowserConfig{Key: "chrome", Name: "Chrome", Kind: types.Chromium, UserDataDir: chromiumDir},
+			wantLen:     1,
+			wantName:    "Chrome",
+			wantProfile: "Default",
+		},
+		{
+			name:        "firefox dispatch",
+			cfg:         types.BrowserConfig{Key: "firefox", Name: "Firefox", Kind: types.Firefox, UserDataDir: firefoxDir},
+			wantLen:     1,
+			wantName:    "Firefox",
+			wantProfile: "abc.default",
+		},
+		{
+			name:    "unknown kind returns error",
+			cfg:     types.BrowserConfig{Key: "unknown", Name: "Unknown", Kind: types.BrowserKind(99)},
+			wantErr: "unknown browser kind",
+		},
+		{
+			name: "empty dir returns empty",
+			cfg:  types.BrowserConfig{Key: "chrome", Name: "Chrome", Kind: types.Chromium, UserDataDir: emptyDir},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			found, err := newBrowsers(tt.cfg)
+			if tt.wantErr != "" {
+				require.Error(t, err)
+				assert.Contains(t, err.Error(), tt.wantErr)
+				return
+			}
+			require.NoError(t, err)
+			require.Len(t, found, tt.wantLen)
+			if tt.wantLen > 0 {
+				assert.Equal(t, tt.wantName, found[0].BrowserName())
+				assert.Equal(t, tt.wantProfile, found[0].ProfileName())
+			}
+		})
+	}
+}
+
+// ---------------------------------------------------------------------------
+// Helpers
+// ---------------------------------------------------------------------------
+
 // assertBrowsers verifies browser names and profiles match expectations (order-independent).
 func assertBrowsers(t *testing.T, browsers []Browser, wantNames, wantProfiles []string) {
 	t.Helper()
@@ -1,6 +1,7 @@
 package chromium

 import (
+	"fmt"
 	"os"
 	"path/filepath"
 	"time"
@@ -16,15 +17,15 @@ import (
 type Browser struct {
 	cfg         types.BrowserConfig
 	profileDir  string                               // absolute path to profile directory
-	retriever   keyretriever.KeyRetriever            // shared across profiles of the same browser
+	retriever   keyretriever.KeyRetriever            // set via SetRetriever after construction
 	sources     map[types.Category][]sourcePath      // Category → candidate paths (priority order)
 	extractors  map[types.Category]categoryExtractor // Category → custom extract function override
 	sourcePaths map[types.Category]resolvedPath      // Category → discovered absolute path
 }

 // NewBrowsers discovers Chromium profiles under cfg.UserDataDir and returns
-// one Browser per profile. Uses ReadDir to find profile directories,
-// then Stat to check which data sources exist in each profile.
+// one Browser per profile. Call SetRetriever on each returned browser before
+// Extract to enable decryption of sensitive data (passwords, cookies, etc.).
 func NewBrowsers(cfg types.BrowserConfig) ([]*Browser, error) {
 	sources := sourcesForKind(cfg.Kind)
 	extractors := extractorsForKind(cfg.Kind)
@@ -34,11 +35,6 @@ func NewBrowsers(cfg types.BrowserConfig) ([]*Browser, error) {
 		return nil, nil
 	}

-	// Create the key retriever once and share it across all profiles.
-	// This avoids repeated keychain password prompts on macOS, where each
-	// profile would otherwise trigger a separate `security` command dialog.
-	retriever := keyretriever.DefaultRetriever(cfg.KeychainPassword)
-
 	var browsers []*Browser
 	for _, profileDir := range profileDirs {
 		sourcePaths := resolveSourcePaths(sources, profileDir)
@@ -48,7 +44,6 @@ func NewBrowsers(cfg types.BrowserConfig) ([]*Browser, error) {
 		browsers = append(browsers, &Browser{
 			cfg:         cfg,
 			profileDir:  profileDir,
-			retriever:   retriever,
 			sources:     sources,
 			extractors:  extractors,
 			sourcePaths: sourcePaths,
@@ -57,6 +52,13 @@ func NewBrowsers(cfg types.BrowserConfig) ([]*Browser, error) {
 	return browsers, nil
 }

+// SetRetriever sets the key retriever used by Extract to obtain the
+// master encryption key. Must be called before Extract if encrypted
+// data (passwords, cookies, credit cards) needs to be decrypted.
+func (b *Browser) SetRetriever(r keyretriever.KeyRetriever) {
+	b.retriever = r
+}
+
 func (b *Browser) BrowserName() string { return b.cfg.Name }
 func (b *Browser) ProfileDir() string  { return b.profileDir }
 func (b *Browser) ProfileName() string {
@@ -120,6 +122,10 @@ func (b *Browser) acquireFiles(session *filemanager.Session, categories []types.
 // The retriever is always called regardless of whether Local State exists,
 // because macOS/Linux retrievers don't need it.
 func (b *Browser) getMasterKey(session *filemanager.Session) ([]byte, error) {
+	if b.retriever == nil {
+		return nil, fmt.Errorf("key retriever not set for %s", b.cfg.Name)
+	}
+
 	// Try to locate and copy Local State (needed on Windows, ignored on macOS/Linux).
 	// Multi-profile layout: Local State is in the parent of profileDir.
 	// Flat layout (Opera): Local State is alongside data files in profileDir.
@@ -8,6 +8,7 @@ import (
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"

+	"github.com/moond4rk/hackbrowserdata/crypto/keyretriever"
 	"github.com/moond4rk/hackbrowserdata/filemanager"
 	"github.com/moond4rk/hackbrowserdata/types"
 )
@@ -438,3 +439,172 @@ func TestLocalStatePath(t *testing.T) {
 		})
 	}
 }
+
+// ---------------------------------------------------------------------------
+// getMasterKey
+// ---------------------------------------------------------------------------
+
+// mockRetriever records the arguments passed to RetrieveKey.
+type mockRetriever struct {
+	storage    string
+	localState string
+	key        []byte
+	err        error
+	called     bool
+}
+
+func (m *mockRetriever) RetrieveKey(storage, localStatePath string) ([]byte, error) {
+	m.called = true
+	m.storage = storage
+	m.localState = localStatePath
+	return m.key, m.err
+}
+
+func TestGetMasterKey(t *testing.T) {
+	// Profile directory without Local State file.
+	dirNoLocalState := t.TempDir()
+	mkFile(dirNoLocalState, "Default", "Preferences")
+	mkFile(dirNoLocalState, "Default", "History")
+
+	tests := []struct {
+		name           string
+		dir            string
+		storage        string
+		retriever      keyretriever.KeyRetriever // nil → don't call SetRetriever
+		wantKey        []byte
+		wantErr        string
+		wantStorage    string
+		wantLocalState bool // whether localStatePath passed to retriever is non-empty
+	}{
+		{
+			name:    "nil retriever returns error",
+			dir:     fixture.chrome,
+			wantErr: "key retriever not set",
+		},
+		{
+			name:           "with Local State passes path to retriever",
+			dir:            fixture.chrome,
+			storage:        "Chrome",
+			retriever:      &mockRetriever{key: []byte("fake-master-key")},
+			wantKey:        []byte("fake-master-key"),
+			wantStorage:    "Chrome",
+			wantLocalState: true,
+		},
+		{
+			name:        "without Local State passes empty path",
+			dir:         dirNoLocalState,
+			storage:     "Chromium",
+			retriever:   &mockRetriever{key: []byte("derived-key")},
+			wantKey:     []byte("derived-key"),
+			wantStorage: "Chromium",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			browsers, err := NewBrowsers(types.BrowserConfig{
+				Name: "Test", Kind: types.Chromium, UserDataDir: tt.dir, Storage: tt.storage,
+			})
+			require.NoError(t, err)
+			require.NotEmpty(t, browsers)
+
+			b := browsers[0]
+			if tt.retriever != nil {
+				b.SetRetriever(tt.retriever)
+			}
+
+			session, err := filemanager.NewSession()
+			require.NoError(t, err)
+			defer session.Cleanup()
+
+			key, err := b.getMasterKey(session)
+			if tt.wantErr != "" {
+				require.Error(t, err)
+				assert.Contains(t, err.Error(), tt.wantErr)
+				return
+			}
+			require.NoError(t, err)
+			assert.Equal(t, tt.wantKey, key)
+
+			mock, ok := tt.retriever.(*mockRetriever)
+			require.True(t, ok)
+			assert.True(t, mock.called)
+			assert.Equal(t, tt.wantStorage, mock.storage)
+			if tt.wantLocalState {
+				assert.NotEmpty(t, mock.localState)
+			} else {
+				assert.Empty(t, mock.localState)
+			}
+		})
+	}
+}
+
+// ---------------------------------------------------------------------------
+// Extract
+// ---------------------------------------------------------------------------
+
+func TestExtract(t *testing.T) {
+	// Shared fixture: profile with a real History database.
+	dir := t.TempDir()
+	mkFile(dir, "Default", "Preferences")
+
+	historyDB := createTestDB(t, "History", urlsSchema,
+		insertURL("https://example.com", "Example", 5, 13350000000000000),
+	)
+	profileDir := filepath.Join(dir, "Default")
+	data, err := os.ReadFile(historyDB)
+	require.NoError(t, err)
+	require.NoError(t, os.WriteFile(filepath.Join(profileDir, "History"), data, 0o644))
+
+	tests := []struct {
+		name          string
+		retriever     keyretriever.KeyRetriever // nil → don't call SetRetriever
+		wantRetriever bool                      // whether retriever should be called
+	}{
+		{
+			name: "without retriever extracts unencrypted data",
+		},
+		{
+			name:          "with mock retriever",
+			retriever:     &mockRetriever{key: []byte("test-key-16bytes")},
+			wantRetriever: true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			browsers, err := NewBrowsers(types.BrowserConfig{
+				Name: "Test", Kind: types.Chromium, UserDataDir: dir, Storage: "Chrome",
+			})
+			require.NoError(t, err)
+			require.Len(t, browsers, 1)
+
+			if tt.retriever != nil {
+				browsers[0].SetRetriever(tt.retriever)
+			}
+
+			result, err := browsers[0].Extract([]types.Category{types.History})
+			require.NoError(t, err)
+			require.NotNil(t, result)
+			require.Len(t, result.Histories, 1)
+			assert.Equal(t, "Example", result.Histories[0].Title)
+
+			if tt.wantRetriever {
+				mock, ok := tt.retriever.(*mockRetriever)
+				require.True(t, ok)
+				assert.True(t, mock.called)
+			}
+		})
+	}
+}
+
+// ---------------------------------------------------------------------------
+// SetRetriever: verify *Browser satisfies the interface used by
+// browser.pickFromConfigs for post-construction retriever injection.
+// ---------------------------------------------------------------------------
+
+func TestSetRetriever_SatisfiesInterface(t *testing.T) {
+	var _ interface {
+		SetRetriever(keyretriever.KeyRetriever)
+	} = (*Browser)(nil)
+}