Yasso/cmd/tools.go

package cmd

import (
	"bufio"
	"bytes"
	"database/sql"
	"fmt"
	"github.com/olekukonko/tablewriter"
	"golang.org/x/net/proxy"
	"log"
	"net"
	"os"
	"strings"
	"time"
)

// socks5代理连接功能

func ConnBySOCKS5()(proxy.Dialer,error){
	// 解析连接过来的socks5字符串
	if strings.ContainsAny(ProxyHost,"@") && strings.Count(ProxyHost,"@") == 1 {
		info := strings.Split(ProxyHost,"@")
		userpass := strings.Split(info[0],":")
		auth := proxy.Auth{User:userpass[0],Password:userpass[1]}
		dialer, err := proxy.SOCKS5("tcp",info[1],&auth,proxy.Direct)
		return dialer,err
	} else {
		if strings.ContainsAny(ProxyHost,":") && strings.Count(ProxyHost,":")==1{
			dialer,err:= proxy.SOCKS5("tcp",ProxyHost,nil,proxy.Direct)
			return dialer,err
		}
	}
	return nil,fmt.Errorf("proxy error")
}

// 返回一个连接

func GetConn(addr string,timeout time.Duration)(net.Conn,error) {
	if ProxyHost != "" {
		dialer,err := ConnBySOCKS5()
		if err != nil {
			return nil,err
		}
		conn, err := dialer.Dial("tcp",addr)
		if err != nil {
			return nil,err
		}
		return conn,nil
	}else{
		return net.DialTimeout("tcp",addr,time.Duration(timeout))
	}
}


func SQLExecute(db *sql.DB,q string)(*Results,error){
	if q == "" {
		return nil, nil
	}
	rows, err := db.Query(q)
	//rows, err := db.Query(q)
	if err != nil {
		return nil, err
	}
	columns, err := rows.Columns()
	if err != nil {
		return nil, err
	}

	var results [][]string
	for rows.Next() {
		rs := make([]sql.NullString, len(columns))
		rsp := make([]interface{}, len(columns))
		for i := range rs {
			rsp[i] = &rs[i]
		}
		if err = rows.Scan(rsp...); err != nil {
			break
		}

		_rs := make([]string, len(columns))
		for i := range rs {
			_rs[i] = rs[i].String
		}
		results = append(results, _rs)
	}
	if closeErr := rows.Close(); closeErr != nil {
		return nil, closeErr
	}
	if err != nil {
		return nil, err
	}
	if err := rows.Err(); err != nil {
		return nil, err
	}
	return &Results{
		Columns: columns,
		Rows:    results,
	}, nil
}

type Results struct {
	Columns []string
	Rows    [][]string
}

func (r *Results) String() string {
	buf := bytes.NewBufferString("")
	table := tablewriter.NewWriter(buf)
	table.SetHeader(r.Columns)
	table.AppendBulk(r.Rows)
	table.Render()
	return buf.String()
}

func SQLshell(db *sql.DB,sqltype string){
	reader := bufio.NewReader(os.Stdin)
	Println(fmt.Sprintf("Welcome to Yasso sql client "))
	for {
		fmt.Printf("Yasso-%s> ",sqltype)
		sqlstr, err := reader.ReadString('\n')
		if err != nil {
			log.Panic("failed to ReadString ", err)
		}
		sqlstr = strings.Trim(sqlstr, "\r\n")
		sqls := []byte(sqlstr)
		if len(sqls) > 6 {
			if string(sqls[:6]) == "select" || string(sqls[:4]) == "show" || string(sqls[:4]) == "desc" {
				//result set sql
				r,err := SQLExecute(db,sqlstr)
				if err != nil {
					Println(fmt.Sprintf("%v",err))
				}
				Println(fmt.Sprintf("%v",r))
			} else {
				//no result set sql
				r,err := SQLExecute(db,sqlstr)
				if err != nil {
					Println(fmt.Sprintf("%v",err))
				}
				Println(fmt.Sprintf("%v",r))
			}
		}
		if sqlstr == "exit" {
			Println("exit sql shell")
			break
		}
	}
}