mirror of
https://github.com/lukaszraczylo/mcp-readitall.git
synced 2026-10-12 00:45:33 +00:00
initial commit
This commit is contained in:
@@ -0,0 +1,327 @@
|
||||
package browser
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/url"
|
||||
"os"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/chromedp/cdproto/cdp"
|
||||
"github.com/chromedp/cdproto/network"
|
||||
"github.com/chromedp/chromedp"
|
||||
)
|
||||
|
||||
// storageState is the on-disk JSON shape for a saved session. It is
|
||||
// intentionally identical to Playwright's storage state so sessions
|
||||
// captured by either tool can be loaded by the other.
|
||||
//
|
||||
// Cookies and localStorage entries are stored as their raw JSON-encoded
|
||||
// representation, not as strongly-typed fields, so the file remains
|
||||
// tolerant of optional / future fields in the upstream format.
|
||||
type storageState struct {
|
||||
Cookies []json.RawMessage `json:"cookies"`
|
||||
Origins []json.RawMessage `json:"origins"`
|
||||
}
|
||||
|
||||
// localStorageEntry is the {name, value} pair shape Playwright (and we)
|
||||
// use to persist per-origin localStorage. Defined at package level so
|
||||
// buildLocalStorageScript can take a named slice type.
|
||||
type localStorageEntry struct {
|
||||
Name string `json:"name"`
|
||||
Value string `json:"value"`
|
||||
}
|
||||
|
||||
// localStorageOrigin groups localStorage entries by origin.
|
||||
type localStorageOrigin struct {
|
||||
Origin string `json:"origin"`
|
||||
LocalStorage []localStorageEntry `json:"localStorage"`
|
||||
}
|
||||
|
||||
// readStorageState parses a storage state file into the in-memory shape.
|
||||
func readStorageState(path string) (*storageState, error) {
|
||||
raw, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("read: %w", err)
|
||||
}
|
||||
var s storageState
|
||||
if err := json.Unmarshal(raw, &s); err != nil {
|
||||
return nil, fmt.Errorf("parse: %w", err)
|
||||
}
|
||||
return &s, nil
|
||||
}
|
||||
|
||||
// ApplyStorageStateCookies sets every cookie in the storage state on the
|
||||
// chromedp context. Cookies are domain-bound and can be set without being
|
||||
// on the target origin, so we apply them pre-navigation.
|
||||
//
|
||||
// Exported so the Browser wrapper can call it from NewContext; not part
|
||||
// of the public API for tool authors.
|
||||
func ApplyStorageStateCookies(ctx context.Context, statePath string) error {
|
||||
state, err := readStorageState(statePath)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if len(state.Cookies) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
cookies, err := decodeCookies(state.Cookies)
|
||||
if err != nil {
|
||||
return fmt.Errorf("decode cookies: %w", err)
|
||||
}
|
||||
if len(cookies) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
return chromedp.Run(ctx, chromedp.ActionFunc(func(ctx context.Context) error {
|
||||
target := chromedp.FromContext(ctx).Target
|
||||
return network.SetCookies(cookies).Do(cdp.WithExecutor(ctx, target))
|
||||
}))
|
||||
}
|
||||
|
||||
// ApplyLocalStorageOnCurrentOrigin sets every localStorage entry that
|
||||
// belongs to the page's current origin. Must be called after the page has
|
||||
// navigated, because localStorage is origin-scoped.
|
||||
//
|
||||
// Entries belonging to other origins are silently skipped: they cannot be
|
||||
// set from the wrong origin anyway.
|
||||
//
|
||||
// Exported so the reader can call it after navigation.
|
||||
func ApplyLocalStorageOnCurrentOrigin(ctx context.Context, statePath string) error {
|
||||
state, err := readStorageState(statePath)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if len(state.Origins) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
var currentURL string
|
||||
if err := chromedp.Run(ctx, chromedp.Location(¤tURL)); err != nil {
|
||||
return fmt.Errorf("read current origin: %w", err)
|
||||
}
|
||||
currentOrigin := normalizeOrigin(currentURL)
|
||||
if currentOrigin == "" {
|
||||
return nil
|
||||
}
|
||||
|
||||
var entries []localStorageEntry
|
||||
for _, raw := range state.Origins {
|
||||
var o localStorageOrigin
|
||||
if err := json.Unmarshal(raw, &o); err != nil {
|
||||
continue
|
||||
}
|
||||
if normalizeOrigin(o.Origin) != currentOrigin {
|
||||
continue
|
||||
}
|
||||
entries = append(entries, o.LocalStorage...)
|
||||
}
|
||||
if len(entries) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
return chromedp.Run(ctx, chromedp.Evaluate(buildLocalStorageScript(entries), nil))
|
||||
}
|
||||
|
||||
// CaptureStorageState reads every cookie set on the current target and
|
||||
// every localStorage entry for the current origin, then marshals them to
|
||||
// the on-disk JSON shape. It is intended to be called by the login tool
|
||||
// after a successful interactive authentication.
|
||||
func CaptureStorageState(ctx context.Context, statePath string) error {
|
||||
// 1) Cookies
|
||||
cookies, err := captureCDPCookies(ctx)
|
||||
if err != nil {
|
||||
return fmt.Errorf("get cookies: %w", err)
|
||||
}
|
||||
|
||||
// 2) Current origin + localStorage
|
||||
var currentURL string
|
||||
if err := chromedp.Run(ctx, chromedp.Location(¤tURL)); err != nil {
|
||||
return fmt.Errorf("read location: %w", err)
|
||||
}
|
||||
origin := normalizeOrigin(currentURL)
|
||||
|
||||
lsRaw, err := captureLocalStorage(ctx, origin)
|
||||
if err != nil {
|
||||
return fmt.Errorf("read localStorage: %w", err)
|
||||
}
|
||||
|
||||
// 3) Build the on-disk shape.
|
||||
cookieMsgs := make([]json.RawMessage, 0, len(cookies))
|
||||
for _, c := range cookies {
|
||||
raw, err := json.Marshal(c)
|
||||
if err != nil {
|
||||
return fmt.Errorf("encode cookie: %w", err)
|
||||
}
|
||||
cookieMsgs = append(cookieMsgs, raw)
|
||||
}
|
||||
|
||||
out := storageState{Cookies: cookieMsgs}
|
||||
if lsRaw != nil {
|
||||
out.Origins = []json.RawMessage{lsRaw}
|
||||
}
|
||||
|
||||
// 4) Write atomically: tmp + rename, mode 0600.
|
||||
body, err := json.MarshalIndent(out, "", " ")
|
||||
if err != nil {
|
||||
return fmt.Errorf("encode state: %w", err)
|
||||
}
|
||||
tmp := statePath + ".tmp"
|
||||
if err := os.WriteFile(tmp, body, 0o600); err != nil {
|
||||
return fmt.Errorf("write tmp: %w", err)
|
||||
}
|
||||
if err := os.Rename(tmp, statePath); err != nil {
|
||||
return fmt.Errorf("rename: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// captureCDPCookies asks the current target for its full cookie list via
|
||||
// the network domain.
|
||||
func captureCDPCookies(ctx context.Context) ([]*network.Cookie, error) {
|
||||
var out []*network.Cookie
|
||||
err := chromedp.Run(ctx, chromedp.ActionFunc(func(ctx context.Context) error {
|
||||
target := chromedp.FromContext(ctx).Target
|
||||
var err error
|
||||
out, err = network.GetCookies().Do(cdp.WithExecutor(ctx, target))
|
||||
return err
|
||||
}))
|
||||
return out, err
|
||||
}
|
||||
|
||||
// captureLocalStorage serialises the current page's localStorage to a
|
||||
// JSON-encoded localStorageOrigin, or returns nil if there is no
|
||||
// scannable origin / no entries.
|
||||
func captureLocalStorage(ctx context.Context, origin string) (json.RawMessage, error) {
|
||||
if origin == "" {
|
||||
return nil, nil
|
||||
}
|
||||
const script = `JSON.stringify(Object.fromEntries(Object.keys(localStorage).map(function(k){return [k, localStorage.getItem(k)];})))`
|
||||
var lsJSON string
|
||||
if err := chromedp.Run(ctx, chromedp.Evaluate(script, &lsJSON)); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if lsJSON == "" || lsJSON == "[]" || lsJSON == "{}" {
|
||||
return nil, nil
|
||||
}
|
||||
var kv [][2]string
|
||||
if err := json.Unmarshal([]byte(lsJSON), &kv); err != nil {
|
||||
return nil, fmt.Errorf("decode localStorage: %w", err)
|
||||
}
|
||||
if len(kv) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
entries := make([]localStorageEntry, 0, len(kv))
|
||||
for _, pair := range kv {
|
||||
entries = append(entries, localStorageEntry{Name: pair[0], Value: pair[1]})
|
||||
}
|
||||
o := localStorageOrigin{Origin: origin, LocalStorage: entries}
|
||||
raw, err := json.Marshal(o)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return raw, nil
|
||||
}
|
||||
|
||||
// decodeCookies turns the on-disk json.RawMessage cookie list into CDP
|
||||
// *network.CookieParam values ready for network.SetCookies. Session
|
||||
// cookies (negative or zero Expires) are left with a nil Expires field,
|
||||
// which CDP interprets as "session cookie".
|
||||
func decodeCookies(raws []json.RawMessage) ([]*network.CookieParam, error) {
|
||||
out := make([]*network.CookieParam, 0, len(raws))
|
||||
for _, raw := range raws {
|
||||
var c onDiskCookie
|
||||
if err := json.Unmarshal(raw, &c); err != nil {
|
||||
return nil, fmt.Errorf("cookie %s: %w", string(raw), err)
|
||||
}
|
||||
out = append(out, onDiskToCookieParam(c))
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// onDiskCookie is the subset of cookie fields we care about when
|
||||
// reading a stored state file. Anything else (priority, source scheme,
|
||||
// partition keys…) is ignored.
|
||||
type onDiskCookie struct {
|
||||
Name string `json:"name"`
|
||||
Value string `json:"value"`
|
||||
Domain string `json:"domain"`
|
||||
Path string `json:"path"`
|
||||
Expires float64 `json:"expires"`
|
||||
HTTPOnly bool `json:"httpOnly"`
|
||||
Secure bool `json:"secure"`
|
||||
SameSite string `json:"sameSite"`
|
||||
}
|
||||
|
||||
// onDiskToCookieParam converts our on-disk shape into a CDP
|
||||
// network.CookieParam. Session cookies (Expires <= 0) get a nil Expires
|
||||
// pointer; persistent cookies get a *cdp.TimeSinceEpoch built from the
|
||||
// Unix timestamp.
|
||||
func onDiskToCookieParam(c onDiskCookie) *network.CookieParam {
|
||||
p := &network.CookieParam{
|
||||
Name: c.Name,
|
||||
Value: c.Value,
|
||||
Domain: c.Domain,
|
||||
Path: c.Path,
|
||||
Secure: c.Secure,
|
||||
HTTPOnly: c.HTTPOnly,
|
||||
SameSite: parseSameSite(c.SameSite),
|
||||
}
|
||||
if c.Expires > 0 {
|
||||
t := cdp.TimeSinceEpoch(time.Unix(int64(c.Expires), 0))
|
||||
p.Expires = &t
|
||||
}
|
||||
return p
|
||||
}
|
||||
|
||||
// parseSameSite translates a Playwright-style SameSite string to a CDP
|
||||
// enum. Unknown / empty values become the safe default (Lax) rather
|
||||
// than None, to match the previous Playwright behaviour.
|
||||
func parseSameSite(s string) network.CookieSameSite {
|
||||
switch strings.ToLower(s) {
|
||||
case "strict":
|
||||
return network.CookieSameSiteStrict
|
||||
case "lax":
|
||||
return network.CookieSameSiteLax
|
||||
case "none":
|
||||
return network.CookieSameSiteNone
|
||||
default:
|
||||
return network.CookieSameSiteLax
|
||||
}
|
||||
}
|
||||
|
||||
// buildLocalStorageScript returns a JS snippet that, when run inside the
|
||||
// page, sets every (k, v) pair in the page's localStorage.
|
||||
func buildLocalStorageScript(entries []localStorageEntry) string {
|
||||
var b strings.Builder
|
||||
b.WriteString("(function(){")
|
||||
for _, e := range entries {
|
||||
b.WriteString("localStorage.setItem(")
|
||||
b.WriteString(jsonString(e.Name))
|
||||
b.WriteString(",")
|
||||
b.WriteString(jsonString(e.Value))
|
||||
b.WriteString(");")
|
||||
}
|
||||
b.WriteString("})();")
|
||||
return b.String()
|
||||
}
|
||||
|
||||
// jsonString produces a JS-safe double-quoted string literal.
|
||||
func jsonString(s string) string {
|
||||
b, _ := json.Marshal(s)
|
||||
return string(b)
|
||||
}
|
||||
|
||||
// normalizeOrigin returns the scheme+host[:port] of a URL, or "" if
|
||||
// parsing fails. Used to match storage state origins against the page's
|
||||
// current location.
|
||||
func normalizeOrigin(raw string) string {
|
||||
u, err := url.Parse(raw)
|
||||
if err != nil || u.Host == "" {
|
||||
return ""
|
||||
}
|
||||
return u.Scheme + "://" + u.Host
|
||||
}
|
||||
Reference in New Issue
Block a user