package cmd

import (
	"Yasso/config"
	"database/sql"
	"fmt"
	_ "github.com/lib/pq"
	"github.com/spf13/cobra"
	"time"
)

var PostgreCmd = &cobra.Command{
	Use: "postgres",
	Short: "PostgreSQL burst module (not support proxy)",
	Run: func(cmd *cobra.Command, args []string) {
		if Hosts == ""{
			_ = cmd.Help()
		}else{
			BrutePostgreByUser()
		}
	},
}

func BrutePostgreByUser(){
	if BrutePort == 0 {
		BrutePort = 5432
	}
	var ips []string
	var err error
	if Hosts != ""{
		ips,err = ResolveIPS(Hosts)
		if err != nil {
			Println(fmt.Sprintf("resolve hosts address failed %v",err))
			return
		}
		if BruteFlag == true {
			users,pass := ReadTextToDic("postgres",UserDic,PassDic)
			Println(Clearln+"[*] Brute Module [postgres]")
			Println(fmt.Sprintf(Clearln + "[*] Have [user:%v] [pass:%v] [request:%v]",len(users),len(pass),len(users) * len(pass) * len(ips)))
			SwitchBurp("postgres",users,pass,ips,BrutePort,Runtime,TimeDuration,"")
		}else{
			Println(Clearln + "[*] May be you want to brute? try to add --crack")
		}
	}
}




func PostgreConn(info config.HostIn,user,pass string)(bool, error){
	var flag = false
	db, err := sql.Open("postgres",fmt.Sprintf("postgres://%v:%v@%v:%v/%v?sslmode=%v",user,pass,info.Host,info.Port,"postgres","disable"))
	if err == nil {
		db.SetConnMaxLifetime(time.Duration(info.TimeOut))
		defer db.Close()
		err = db.Ping()
		if err == nil {
			flag = true
		}
	}
	return flag,err
}