126 lines
2.9 KiB
Go
126 lines
2.9 KiB
Go
package geminiweb
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"os"
|
|
"path/filepath"
|
|
"time"
|
|
|
|
"github.com/go-rod/rod"
|
|
"github.com/go-rod/rod/lib/launcher"
|
|
"github.com/go-rod/rod/lib/proto"
|
|
)
|
|
|
|
type Browser struct {
|
|
browser *rod.Browser
|
|
visible bool
|
|
}
|
|
|
|
func NewBrowser(visible bool) (*Browser, error) {
|
|
l := launcher.New()
|
|
if visible {
|
|
l = l.Headless(false)
|
|
} else {
|
|
l = l.Headless(true)
|
|
}
|
|
|
|
url, err := l.Launch()
|
|
if err != nil {
|
|
return nil, fmt.Errorf("failed to launch browser: %w", err)
|
|
}
|
|
|
|
b := rod.New().ControlURL(url)
|
|
if err := b.Connect(); err != nil {
|
|
return nil, fmt.Errorf("failed to connect browser: %w", err)
|
|
}
|
|
|
|
return &Browser{browser: b, visible: visible}, nil
|
|
}
|
|
|
|
func (b *Browser) Close() error {
|
|
if b.browser != nil {
|
|
return b.browser.Close()
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (b *Browser) NewPage() (*rod.Page, error) {
|
|
return b.browser.Page(proto.TargetCreateTarget{URL: "about:blank"})
|
|
}
|
|
|
|
type Cookie 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"`
|
|
}
|
|
|
|
func LoadCookiesFromFile(cookieFile string) ([]Cookie, error) {
|
|
data, err := os.ReadFile(cookieFile)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("failed to read cookies: %w", err)
|
|
}
|
|
|
|
var cookies []Cookie
|
|
if err := json.Unmarshal(data, &cookies); err != nil {
|
|
return nil, fmt.Errorf("failed to parse cookies: %w", err)
|
|
}
|
|
|
|
return cookies, nil
|
|
}
|
|
|
|
func SaveCookiesToFile(cookies []Cookie, cookieFile string) error {
|
|
data, err := json.MarshalIndent(cookies, "", " ")
|
|
if err != nil {
|
|
return fmt.Errorf("failed to marshal cookies: %w", err)
|
|
}
|
|
|
|
dir := filepath.Dir(cookieFile)
|
|
if err := os.MkdirAll(dir, 0755); err != nil {
|
|
return fmt.Errorf("failed to create cookie dir: %w", err)
|
|
}
|
|
|
|
if err := os.WriteFile(cookieFile, data, 0644); err != nil {
|
|
return fmt.Errorf("failed to write cookies: %w", err)
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
func SetCookiesOnPage(page *rod.Page, cookies []Cookie) error {
|
|
var protoCookies []*proto.NetworkCookieParam
|
|
for _, c := range cookies {
|
|
p := &proto.NetworkCookieParam{
|
|
Name: c.Name,
|
|
Value: c.Value,
|
|
Domain: c.Domain,
|
|
Path: c.Path,
|
|
HTTPOnly: c.HTTPOnly,
|
|
Secure: c.Secure,
|
|
}
|
|
if c.Expires > 0 {
|
|
exp := proto.TimeSinceEpoch(c.Expires)
|
|
p.Expires = exp
|
|
}
|
|
protoCookies = append(protoCookies, p)
|
|
}
|
|
return page.SetCookies(protoCookies)
|
|
}
|
|
|
|
func WaitForElement(page *rod.Page, selector string, timeout time.Duration) (*rod.Element, error) {
|
|
ctx, cancel := context.WithTimeout(context.Background(), timeout)
|
|
defer cancel()
|
|
return page.Context(ctx).Element(selector)
|
|
}
|
|
|
|
func WaitForElements(page *rod.Page, selector string, timeout time.Duration) (rod.Elements, error) {
|
|
ctx, cancel := context.WithTimeout(context.Background(), timeout)
|
|
defer cancel()
|
|
return page.Context(ctx).Elements(selector)
|
|
}
|