browserSpider/program/spider/action.go

/**
 * 浏览器行为封装
 * 基础动作
 */
package main

import (
	"context"
	"errors"
	"strconv"
	"strings"
	"time"

	"github.com/chromedp/chromedp"
	"github.com/yuin/gopher-lua"
)

const (
	selector_type_id        = 0
	selector_type_query     = 1
	selector_type_search    = 2
	selector_type_jspath    = 3
	selector_type_query_all = 4

	execute_return_type_string = 0
	execute_return_type_list   = 1
	execute_return_type_table  = 2
)

// findTab 根据标题、url找tab
func (b *Browser) findTabContext(tabTitle, tabUrl string, timeoutInt64 int64) (ctx context.Context, err error) {
	if timeoutInt64 == 0 {
		timeoutInt64 = 5000
	}
	timeout := time.Duration(timeoutInt64) * time.Millisecond
	if tabTitle == "*" && tabUrl == "*" {
		return b.Ctx, nil
	} else if tabTitle == "" && tabUrl == "" {
		ctx, _ = context.WithTimeout(b.Ctx, timeout)
		return ctx, nil
	} else {
		ts, err := chromedp.Targets(b.Ctx)
		if err != nil {
			return nil, err
		}
		for _, t := range ts {
			if (tabTitle != "" && strings.Contains(t.Title, tabTitle)) ||
				(tabUrl != "" && strings.Contains(t.URL, tabUrl)) {
				// log.Printf("find tab param<title,url>: %s %s found %s %s", tabTitle, tabUrl,
				// 	t.Title, t.URL)
				newCtx, _ := chromedp.NewContext(b.Ctx, chromedp.WithTargetID(t.TargetID))
				ctx, _ = context.WithTimeout(newCtx, timeout)
				return ctx, nil
			}
		}
	}
	return nil, errors.New("can't find tab")
}

// BindLuaState
func (b *Browser) BindLuaState(state *lua.LState) {
	//执行暂停
	state.SetGlobal("browser_sleep", state.NewFunction(func(l *lua.LState) int {
		timeout := l.ToInt64(-1)
		if timeout == 0 {
			timeout = 5
		}
		time.Sleep(time.Duration(timeout) * time.Millisecond)
		return 0
	}))
	//关闭tabl页
	state.SetGlobal("browser_closetabs", state.NewFunction(func(l *lua.LState) int {
		tabTitle := l.ToString(-3)
		tabUrl := l.ToString(-2)
		timeout := l.ToInt64(-1)
		if timeout == 0 {
			timeout = 5
		}
		b.CloseTabs(tabTitle, tabUrl, timeout)
		return 0
	}))
	//注册打开地址
	state.SetGlobal("browser_navagite", state.NewFunction(func(l *lua.LState) int {
		tabTitle := l.ToString(-5)
		tabUrl := l.ToString(-4)
		isNewTab := l.ToBool(-3)
		timeout := l.ToInt64(-2)
		targetUrl := l.ToString(-1)
		if err := b.Navigate(tabTitle, tabUrl, isNewTab, targetUrl, timeout); err != nil {
			l.Push(lua.LString(err.Error()))
		} else {
			l.Push(lua.LString("ok"))
		}
		return 1
	}))
	//执行浏览器端js
	state.SetGlobal("browser_executejs", state.NewFunction(func(l *lua.LState) int {
		tabTitle := l.ToString(-5)
		tabUrl := l.ToString(-4)
		timeout := l.ToInt64(-3)
		returnType := l.ToInt(-2)
		script := l.ToString(-1)

		switch returnType {
		case execute_return_type_string: //返回string
			var ret string
			if err := b.ExecuteJS(tabTitle, tabUrl, script, &ret, timeout); err == nil {
				l.Push(lua.LString("ok"))
				l.Push(lua.LString(ret))
			} else {
				l.Push(lua.LString("err"))
				l.Push(lua.LString(err.Error()))
			}
		case execute_return_type_list: //返回list
			var ret = make([]interface{}, 0, 0)
			var tmp = make(map[string]interface{})
			if err := b.ExecuteJS(tabTitle, tabUrl, script, &ret, timeout); err == nil {
				for i, v := range ret {
					tmp[strconv.Itoa(i)] = v
				}
				l.Push(lua.LString("ok"))
				l.Push(MapToTable(tmp))
			} else {
				l.Push(lua.LString("err"))
				l.Push(lua.LString(err.Error()))
			}
		case execute_return_type_table: //返回table
			var ret = make(map[string]interface{})
			if err := b.ExecuteJS(tabTitle, tabUrl, script, &ret, timeout); err == nil {
				l.Push(lua.LString("ok"))
				l.Push(MapToTable(ret))
			} else {
				l.Push(lua.LString("err"))
				l.Push(lua.LString(err.Error()))
			}
		}
		return 2
	}))
	//按键
	state.SetGlobal("browser_keysend", state.NewFunction(func(l *lua.LState) int {
		tabTitle := l.ToString(-6)
		tabUrl := l.ToString(-5)
		timeout := l.ToInt64(-4)

		selectorType := l.ToInt(-3)
		selector := l.ToString(-2)
		words := l.ToString(-1)

		err := b.KeySend(tabTitle, tabUrl, selector, words, selectorType, timeout)
		if err != nil {
			l.Push(lua.LString(err.Error()))
		} else {
			l.Push(lua.LString("ok"))
		}
		return 1
	}))
	//点击
	state.SetGlobal("browser_click", state.NewFunction(func(l *lua.LState) int {
		tabTitle := l.ToString(-5)
		tabUrl := l.ToString(-4)
		timeout := l.ToInt64(-3)
		selectorType := l.ToInt(-2)
		selector := l.ToString(-1)

		err := b.Click(tabTitle, tabUrl, selector, selectorType, timeout)
		if err != nil {
			l.Push(lua.LString(err.Error()))
		} else {
			l.Push(lua.LString("ok"))
		}
		return 1
	}))
	//browser_history_back
	state.SetGlobal("browser_history_back", state.NewFunction(func(l *lua.LState) int {
		tabTitle := l.ToString(-3)
		tabUrl := l.ToString(-2)
		timeout := l.ToInt64(-1)

		err := b.GoHistoryBack(tabTitle, tabUrl, timeout)
		if err != nil {
			l.Push(lua.LString(err.Error()))
		} else {
			l.Push(lua.LString("ok"))
		}
		return 1
	}))
	state.SetGlobal("browser_waitvisible", state.NewFunction(func(l *lua.LState) int {
		tabTitle := l.ToString(-5)
		tabUrl := l.ToString(-4)
		timeout := l.ToInt64(-3)
		selectorType := l.ToInt(-2)
		selector := l.ToString(-1)

		err := b.WaitVisible(tabTitle, tabUrl, selector, selectorType, timeout)
		if err != nil {
			l.Push(lua.LString(err.Error()))
		} else {
			l.Push(lua.LString("ok"))
		}
		return 1
	}))
	//点击
	state.SetGlobal("browser_downloadfile", state.NewFunction(func(l *lua.LState) int {
		tabTitle := l.ToString(-7)
		tabUrl := l.ToString(-6)
		timeout := l.ToInt64(-5)
		selectorType := l.ToInt(-4)
		selector := l.ToString(-3)
		filename := l.ToString(-2)
		save2dir := l.ToString(-1)

		err := b.DownloadFile(tabTitle, tabUrl, timeout, selector, selectorType, filename, save2dir)
		if err != nil {
			l.Push(lua.LString(err.Error()))
		} else {
			l.Push(lua.LString("ok"))
		}
		return 1
	}))

	//关闭tabl页
	state.SetGlobal("browser_closetabs_without", state.NewFunction(func(l *lua.LState) int {
		tabTitle := l.ToString(-3)
		tabUrl := l.ToString(-2)
		timeout := l.ToInt64(-1)
		if timeout == 0 {
			timeout = 5
		}
		b.CloseTabsWithout(tabTitle, tabUrl, timeout)
		return 0
	}))
	//browser_screenshot 网页局部截图
	state.SetGlobal("browser_screenshot", state.NewFunction(func(l *lua.LState) int {
		tabTitle := l.ToString(-6)
		tabUrl := l.ToString(-5)
		timeout := l.ToInt64(-4)
		selectorType := l.ToInt(-3)
		selector := l.ToString(-2)
		filename := l.ToString(-1)
		if timeout == 0 {
			timeout = 5
		}

		if err := b.Screenshot(tabTitle, tabUrl, timeout, selectorType, selector, filename); err != nil {
			l.Push(lua.LString(err.Error()))
		} else {
			l.Push(lua.LString("ok"))
		}
		return 1
	}))
	//browser_print2pdf 整个网页生成pdf
	state.SetGlobal("browser_print2pdf", state.NewFunction(func(l *lua.LState) int {
		tabTitle := l.ToString(-4)
		tabUrl := l.ToString(-3)
		timeout := l.ToInt64(-2)
		filename := l.ToString(-1)
		if timeout == 0 {
			timeout = 5
		}

		if err := b.PrintToPDF(tabTitle, tabUrl, timeout, filename); err != nil {
			l.Push(lua.LString(err.Error()))
		} else {
			l.Push(lua.LString("ok"))
		}
		return 1
	}))
	state.SetGlobal("browser_tabs", state.NewFunction(func(l *lua.LState) int {
		tabTitle := l.ToString(-3)
		tabUrl := l.ToString(-2)
		timeout := l.ToInt64(-1)
		if timeout == 0 {
			timeout = 500
		}
		var tmp = make(map[string]interface{})
		tabs, err := b.GetBrowserTabs(tabTitle, tabUrl, timeout)
		//fmt.Println(err, tabs)
		if err == nil {
			for i, v := range tabs {
				tmp[strconv.Itoa(i)] = v
			}
			l.Push(lua.LString("ok"))
			l.Push(MapToTable(tmp))
		} else {
			l.Push(lua.LString("err"))
			l.Push(lua.LString(err.Error()))
		}
		return 2
	}))

	state.SetGlobal("browser_send_img_chatbot", state.NewFunction(func(l *lua.LState) int {
		mentioned := l.ToString(-3)
		uri := l.ToString(-2)
		img := l.ToString(-1)
		err := SendImage2ChatBot(uri, img, mentioned)
		if err != nil {
			l.Push(lua.LString("err"))
			l.Push(lua.LString(err.Error()))
		} else {
			l.Push(lua.LString("ok"))
			l.Push(lua.LString("ok"))
		}
		return 2
	}))
	state.SetGlobal("browser_send_text_chatbot", state.NewFunction(func(l *lua.LState) int {
		mentioned := l.ToString(-3)
		uri := l.ToString(-2)
		text := l.ToString(-1)
		err := SendImage2ChatBot(uri, text, mentioned)
		if err != nil {
			l.Push(lua.LString("err"))
			l.Push(lua.LString(err.Error()))
		} else {
			l.Push(lua.LString("ok"))
			l.Push(lua.LString("ok"))
		}
		return 2
	}))

	state.SetGlobal("browser_reload", state.NewFunction(func(l *lua.LState) int {
		tabTitle := l.ToString(-3)
		tabUrl := l.ToString(-2)
		timeout := l.ToInt64(-1)
		err := b.Reload(tabTitle, tabUrl, timeout)
		if err != nil {
			l.Push(lua.LString("err"))
		} else {
			l.Push(lua.LString("ok"))
		}
		return 1
	}))
	state.SetGlobal("browser_back", state.NewFunction(func(l *lua.LState) int {
		tabTitle := l.ToString(-3)
		tabUrl := l.ToString(-2)
		timeout := l.ToInt64(-1)
		err := b.Back(tabTitle, tabUrl, timeout)
		if err != nil {
			l.Push(lua.LString("err"))
		} else {
			l.Push(lua.LString("ok"))
		}
		return 1
	}))
	state.SetGlobal("browser_forward", state.NewFunction(func(l *lua.LState) int {
		tabTitle := l.ToString(-3)
		tabUrl := l.ToString(-2)
		timeout := l.ToInt64(-1)
		err := b.Forward(tabTitle, tabUrl, timeout)
		if err != nil {
			l.Push(lua.LString("err"))
		} else {
			l.Push(lua.LString("ok"))
		}
		return 1
	}))
}