intimate/utils.go

package intimate

import (
	"crypto/md5"
	"database/sql"
	"fmt"
	"log"
	"net"
	"os"
	"os/exec"
	"os/signal"
	"strconv"
	"strings"
	"sync/atomic"
	"syscall"
	"time"

	"github.com/tebeka/selenium"
	"github.com/tebeka/selenium/chrome"
)

var zeroTime time.Time

func init() {

	tm, err := time.Parse("15:04:05", "0:00:00")
	if err != nil {
		log.Println(err)
	}
	zeroTime = tm

}

// GetUpdateTimeNow 获取马上更新时间.  与第一次连用
func GetUpdateTimeNow() *sql.NullTime {
	return &sql.NullTime{Time: time.Now().Add(-time.Hour * 100000), Valid: true}
}

func GetUrlHash(urlstr string) string {
	return fmt.Sprintf("%x", md5.Sum([]byte(urlstr)))
}

// ParseNumber 去逗号解析数字
func ParseNumber(num string) (int64, error) {
	num = strings.Trim(num, " ")
	num = strings.ReplaceAll(num, ",", "")
	return strconv.ParseInt(num, 10, 64)
}

// ParseNumberEx 解析带字符的数字
func ParseNumberEx(num string) (float64, error) {
	num = strings.Trim(num, " ")
	num = strings.ReplaceAll(num, ",", "")
	last := num[len(num)-1]
	factor := 1.0
	switch {
	case last == 'k' || last == 'K':
		factor = 1000.0
		num = num[0 : len(num)-1]
	case last == 'm' || last == 'M':
		factor = 1000000.0
		num = num[0 : len(num)-1]
	}
	i, err := strconv.ParseFloat(num, 64)
	if err != nil {
		return 0, err
	}

	return i * factor, nil
}

// ParseDuration time to duration eg: 1:40:00 -> time.Duration
func ParseDuration(dt string) (time.Duration, error) {

	var parse []byte = []byte("00:00:00")

	j := len(parse) - 1
	for i := len(dt) - 1; i >= 0; i-- {
		c := dt[i]
		if c != ':' {
			parse[j] = dt[i]
		} else {
			for parse[j] != ':' {
				j--
			}
		}
		j--
	}

	tdt, err := time.Parse("15:04:05", string(parse))
	if err != nil {
		return time.Duration(0), err
	}
	return tdt.Sub(zeroTime), nil
}

type AutoCloseDriver struct {
	Webdriver selenium.WebDriver
	Port      int
}

func (adriver *AutoCloseDriver) Close() {

	data, err := exec.Command("/bin/bash", "-c", fmt.Sprintf(`pgrep -f "port=%d"`, adriver.Port)).Output()
	if err != nil {
		log.Println(err)
		log.Println(string(data))
		return
	}
	// log.Println(string(data))

	killshell := fmt.Sprintf("pkill -9 -P %s", data)
	err = exec.Command("/bin/bash", "-c", killshell).Run()
	if err != nil {
		log.Println(err)
		return
	}

	err = exec.Command("/bin/bash", "-c", fmt.Sprintf("kill -9 %s", data)).Run()
	if err != nil {
		log.Println(err)
		return
	}
}

func GetChromeDriver() *AutoCloseDriver {

	port := GetFreePort()

	var err error
	caps := selenium.Capabilities{"browserName": "chrome"}

	chromecaps := chrome.Capabilities{}
	for _, epath := range []string{"../../../crx/myblock.crx", "../../crx/myblock.crx"} {
		_, err := os.Stat(epath)
		if err == nil {

			err := chromecaps.AddExtension(epath)
			if err != nil {
				panic(err)
			}
			break
		}
	}

	if proxy := os.Getenv("chrome_proxy"); proxy != "" {
		log.Println("proxy-server", proxy)
		chromecaps.Args = append(chromecaps.Args, "--proxy-server="+proxy)
	}

	if proxy := os.Getenv("pac_proxy"); proxy != "" {
		log.Println("--proxy-pac-url=" + proxy)
		chromecaps.Args = append(chromecaps.Args, "--proxy-pac-url="+proxy)
	}

	// chromecaps.Args = append(chromecaps.Args, "--proxy-pac-url=http://127.0.0.1:1081/pac")
	chromecaps.Args = append(chromecaps.Args, "--disk-cache-dir=/tmp/chromedriver-cache")
	chromecaps.Args = append(chromecaps.Args, "--disable-gpu", "--disable-images", "--start-maximized", "--disable-infobars")
	// chromecaps.Args = append(chromecaps.Args, "--headless")
	chromecaps.Args = append(chromecaps.Args, "--no-sandbox")
	chromecaps.Args = append(chromecaps.Args, "--disable-dev-shm-usage", "--mute-audio", "--safebrowsing-disable-auto-update")

	chromecaps.ExcludeSwitches = append(chromecaps.ExcludeSwitches, "enable-automation")
	caps.AddChrome(chromecaps)

	_, err = selenium.NewChromeDriverService("/usr/bin/chromedriver", port)
	if err != nil {
		panic(err)
	}

	wd, err := selenium.NewRemote(caps, fmt.Sprintf("http://localhost:%d/wd/hub", port))
	if err != nil {
		panic(err)
	}

	adriver := &AutoCloseDriver{}
	adriver.Port = port
	adriver.Webdriver = wd

	wd.ExecuteScript("windows.navigator.webdriver = undefined", nil)
	if err != nil {
		panic(err)
	}

	return adriver
}

// PerfectShutdown 完美关闭程序
type PerfectShutdown struct {
	loop int32
}

// NewPerfectShutdown 创建完美关闭程序
func NewPerfectShutdown() *PerfectShutdown {
	ps := &PerfectShutdown{}
	ps.loop = 1

	go func() {
		signalchan := make(chan os.Signal)
		signal.Notify(signalchan, syscall.SIGINT, syscall.SIGKILL, syscall.SIGQUIT, syscall.SIGTERM, syscall.SIGSTOP)
		log.Println("accept stop command:", <-signalchan)
		atomic.StoreInt32(&ps.loop, 0)
	}()

	return ps
}

// IsClose 判断是否要关闭
func (ps *PerfectShutdown) IsClose() bool {
	return atomic.LoadInt32(&ps.loop) == 0
}

// Wait 判断是否要关闭
func (ps *PerfectShutdown) Wait(tm time.Duration) bool {
	now := time.Now()
	for time.Now().Sub(now) <= tm {
		if ps.IsClose() {
			return false
		}
		time.Sleep(time.Second)
	}
	return true
}

type Counter struct {
	dcount   int
	count    int
	maxLimit int
	minLimit int

	minobj       []interface{}
	maxobj       []interface{}
	maxLimitToDo func(obj ...interface{}) error
	minLimitToDo func(obj ...interface{}) error
}

func NewCounter() *Counter {
	c := &Counter{}
	return c
}

// SetDefault 设置默认值
func (c *Counter) SetDefault(n int) {
	c.dcount = n
}

// Reset 最置count为defaultCount值
func (c *Counter) Reset() {
	c.count = c.dcount
}

// SetCount 设置count到最大值的时候执行do函数
func (c *Counter) SetCount(count int) {
	c.count = count
}

// GetCount 设置count到最大值的时候执行do函数
func (c *Counter) GetCount() int {
	return c.count
}

// SetMinLimit 设置最小限制
func (c *Counter) SetMinLimit(n int) {
	c.minLimit = n
}

// SetMaxLimit 设置最大限制
func (c *Counter) SetMaxLimit(n int) {
	c.maxLimit = n
}

// SetMaxToDo 设置count到最大值的时候执行do函数
func (c *Counter) SetMaxToDo(do func(obj ...interface{}) error, obj ...interface{}) {
	c.maxLimitToDo = do
	c.maxobj = obj
}

// SetMinToDo 设置count到最小值的时候执行do函数
func (c *Counter) SetMinToDo(do func(obj ...interface{}) error, obj ...interface{}) {
	c.minLimitToDo = do
	c.minobj = obj
}

// AddWithReset 操作 count 默认值为0, 当触发限制时, 重置为默认值
func (c *Counter) AddWithReset(n int) error {
	c.count += n
	if c.maxLimitToDo != nil {
		if c.count >= c.maxLimit {
			defer c.Reset()
			return c.maxLimitToDo(c.maxobj...)
		}
	}
	if c.minLimitToDo != nil {
		if c.count <= c.minLimit {
			defer c.Reset()
			return c.minLimitToDo(c.minobj...)
		}
	}
	return nil
}

// Add 操作 count 默认值为0
func (c *Counter) Add(n int) error {
	c.count += n
	if c.maxLimitToDo != nil {
		if c.count >= c.maxLimit {
			return c.maxLimitToDo(c.maxobj...)
		}
	}
	if c.minLimitToDo != nil {
		if c.count <= c.minLimit {
			return c.minLimitToDo(c.minobj...)
		}
	}
	return nil
}

type WaitFor struct {
	WebDriver selenium.WebDriver
}

func NewWaitFor(wd selenium.WebDriver) *WaitFor {
	return &WaitFor{WebDriver: wd}
}

func (wf *WaitFor) Default(xpath string, do func(elements ...selenium.WebElement) bool) error {
	return wf.WaitWithTimeout(xpath, 15*time.Second, do)
}

func (wf *WaitFor) WaitWithTimeout(xpath string, timeout time.Duration, do func(elements ...selenium.WebElement) bool) error {
	return wf.WebDriver.WaitWithTimeout(func(wd selenium.WebDriver) (bool, error) {
		elements, err := wd.FindElements(selenium.ByXPATH, xpath)
		if err != nil {
			log.Println(err)
			return false, err
		}

		if len(elements) > 0 {
			if do == nil {
				return true, nil
			}
			if do(elements...) {
				return true, nil
			}
		}
		return false, nil

	}, timeout)
}

func GetFreePort() int {
	addr, err := net.ResolveTCPAddr("tcp", "localhost:0")
	if err != nil {
		panic(err)
	}

	l, err := net.ListenTCP("tcp", addr)
	if err != nil {
		panic(err)
	}
	defer l.Close()
	return l.Addr().(*net.TCPAddr).Port
}