package main

import (
	"flag"
	"fmt"
	"github.com/miekg/dns"
	"html/template"
	"log"
	"math/rand"
	"net"
	"net/http"
	"os"
	"regexp"
	"strings"
	"sync"
	"time"
)

const (
	DEFAULT_PORT    = 53
	DEFAULT_TIMEOUT = 20 * time.Second
	DEFAULT_PROTO   = "udp"
)

var (
	LOG       *log.Logger
	FAKE_IP   net.IP
	PROTO     string
	TIMEOUT   time.Duration
	DEBUG     bool
	TCP_REGEX *regexp.Regexp

	DEFAULT_SERVERS = []string{
		// "8.8.8.8",      // google
		// "8.8.4.4",      // google
		// "208.67.222.222", // OpenDNS
		// "208.67.220.220", // OpenDNS
		"114.114.114.114", // 114dns
		"180.76.76.76",    // baidu
		"223.5.5.5",       // ali
		"223.6.6.6",       // ali
		"119.29.29.29",    // tencent dnspod
	}

	LOGIN_USER = make([]int64, 0, 8)
	USER_MUTEX = &sync.RWMutex{}

	REDIRECT_MODE int
)

func ipv4ToInt64(ip net.IP) int64 {
	if DEBUG {
		fmt.Println("ip to int64: ", ip)
	}
	ip4 := ip.To4()
	if len(ip4) == net.IPv4len {
		n := (int64(ip4[0]) << 24) | (int64(ip4[1]) << 16) | (int64(ip4[2]) << 8) | int64(ip[3])
		return n
	}
	return 0
}
func userExist(ip int64) bool {
	USER_MUTEX.RLock()
	found := false
	for i := 0; i < len(LOGIN_USER); i++ {
		if LOGIN_USER[i] == ip {
			if DEBUG {
				fmt.Println("User Exist: ", ip)
			}
			found = true
			break
		}
	}
	USER_MUTEX.RUnlock()
	return found
}
func addUser(ip int64) {
	if !userExist(ip) {
		if DEBUG {
			fmt.Println("Add User:", ip)
		}
		USER_MUTEX.Lock()
		LOGIN_USER = append(LOGIN_USER, ip)
		USER_MUTEX.Unlock()
	}
}

func init() {
	LOG = log.New(os.Stderr, "[DNS PROXY] ", log.LstdFlags)

	var proto, pattern, fakeIP string
	var timeout int

	flag.BoolVar(&DEBUG, "debug", false, "调试开关")
	flag.StringVar(&proto, "protocol", "udp", "udp|tcp 用udp还是tcp连接远程dns服务器")
	flag.IntVar(&timeout, "timeout", 0, "timeout 等待远程dns服务器的超时时间, 单位秒")
	flag.StringVar(&pattern, "regex", ".*", "需要过滤域名的正则表达式,默认所有域名都返回假ip")
	flag.StringVar(&fakeIP, "fake_ip", "", "dns查询返回的假ip地址")
	flag.IntVar(&REDIRECT_MODE, "redirect_mode", 0, "网页重定向的方式 0到4数字")
	flag.Parse()

	switch proto {
	case "tcp":
		PROTO = "tcp"
	case "udp":
		PROTO = "udp"
	default:
		PROTO = DEFAULT_PROTO
	}

	if timeout > 5 {
		TIMEOUT = time.Duration(timeout) * time.Second
	} else {
		TIMEOUT = DEFAULT_TIMEOUT
	}

	if len(pattern) > 0 {
		if re, err := regexp.Compile(pattern); err != nil {
			LOG.Fatalf("Compiling pattern [%s] was %s", pattern, err)
		} else {
			TCP_REGEX = re
		}
	}
	FAKE_IP = net.ParseIP(fakeIP)

	LOG.Printf("Use %s protocol to connect to the remote DNS server", PROTO)
	LOG.Printf("Timeout duration: %s", TIMEOUT)
	if TCP_REGEX != nil {
		LOG.Printf("Compiling tcp regex pattern [%s]", TCP_REGEX)
	}
	LOG.Printf("Use fake IP %s", fakeIP)
	LOG.Println("Redirect mode: ", REDIRECT_MODE)
}

type Proxy struct {
	Servers []string
}

func (p Proxy) ServeDNS(w dns.ResponseWriter, req *dns.Msg) {
	if TCP_REGEX != nil && FAKE_IP != nil {
		q := req.Question[0]

		isknownUser := false
		ip, _, err := net.SplitHostPort(w.RemoteAddr().String())
		if DEBUG {
			fmt.Println("dns client ip: ", ip)
		}
		if err == nil {
			userIP := net.ParseIP(ip)
			if userIP != nil {
				isknownUser = userExist(ipv4ToInt64(userIP))
			}
		}
		if TCP_REGEX.MatchString(q.Name) && !isknownUser {
			if DEBUG {
				LOG.Printf("DEBUG return a fake ip: %s => %s", q.Name, FAKE_IP)
			}
			m := new(dns.Msg)
			m.SetReply(req)
			m.Authoritative = true
			m.RecursionAvailable = true
			m.Answer = make([]dns.RR, 0, 1)
			m.Answer = append(m.Answer,
				&dns.A{
					Hdr: dns.RR_Header{
						Name:   q.Name,
						Rrtype: dns.TypeA,
						Class:  dns.ClassINET,
						Ttl:    0,
					},
					A: FAKE_IP.To4(),
				})
			w.WriteMsg(m)
			return
		}
	}

	c := new(dns.Client)

	c.Net = PROTO
	// if TCP_REGEX != nil {
	//	for _, q := range r.Question {
	//		if TCP_REGEX.MatchString(q.Name) {
	//			LOG.Printf("Tcp proto regex match: %s", q.Name)
	//			c.Net = "tcp"
	//			break
	//		}
	//	}
	// }

	c.ReadTimeout = TIMEOUT
	c.WriteTimeout = TIMEOUT

	if rs, _, err := c.Exchange(req, p.Server()); err == nil {
		if DEBUG {
			LOG.Printf("DEBUG %s\n%v", w.RemoteAddr(), rs)
		}
		w.WriteMsg(rs)
	} else {
		m := new(dns.Msg)
		m.SetRcode(req, dns.RcodeServerFailure)
		w.WriteMsg(m)
		LOG.Printf("%s %s", w.RemoteAddr(), err)
	}
}

func (p Proxy) Server() string {
	sl := len(p.Servers)
	if sl > 0 {
		i := rand.Intn(sl)
		s := p.Servers[i]
		if strings.Index(s, ":") == -1 {
			s = fmt.Sprintf("%s:%d", s, DEFAULT_PORT)
		}
		return s
	}
	return "8.8.8.8:53"
}

// ----------------------------------------------------------------------------

type PortalServer struct{}

func (h *PortalServer) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	var urlPath = r.URL.Path
	fmt.Println("query url:", urlPath)
	if urlPath == "/login.go" {
		if len(r.FormValue("UserName")) != 0 {
			if len(r.FormValue("WISPrVersion")) != 0 ||
				len(r.FormValue("FNAME")) != 0 ||
				len(r.FormValue("OriginatingServer")) != 0 {
				fmt.Println("用户使用wispr方式登录")
			}

			fmt.Fprint(w, "<HTML><HEAD><TITLE>Success</TITLE></HEAD><BODY>"+r.FormValue("id")+" 登录成功"+"</BODY> </HTML>")
			ip, _, err := net.SplitHostPort(r.RemoteAddr)
			if err != nil {
				fmt.Println("Failed to SplitHostPort:", r.RemoteAddr)
				return
			}
			userIP := net.ParseIP(ip)
			if userIP == nil {
				fmt.Println("Failed to ParseIP:", ip)
				return
			}
			addUser(ipv4ToInt64(userIP))
			// 页面返回后，iPhone很快就下发探测包。我们慢点回复确保它探测的时候AddUser
			// 规则已经其作用了。
			time.Sleep(500 * time.Millisecond)
			if DEBUG {
				fmt.Println("user has submit id to login.go:", ip)
			}
		} else {
			fmt.Fprint(w,
				`<HTML>
<HEAD>
<TITLE>portal</TITLE>
</HEAD>
<BODY>
<form action="login.go" method="post">
<table border="0">
<tr>
<td>用户名:</td> <td> <input type="text" name="UserName" /> </td>
</tr>
<tr>
<td>密码:</td> <td><input type="text" name="Password" /> </td>
</tr>
<tr>
<td></td>
<td><input type="submit" value="提交" /> </td>
<tr>
</form>
</BODY>
</HTML> `)
		}
		return
	}

	switch REDIRECT_MODE {
	default:
		// case 0: // 直接输出
		fmt.Fprint(w,
			`<HTML>
<HEAD>
<TITLE>portal</TITLE>
</HEAD>
<BODY>
<form action="login.go" method="post">
<table border="0">
<tr>
<td>用户名:</td> <td> <input type="text" name="UserName" /> </td>
</tr>
<tr>
<td>密码:</td> <td><input type="text" name="Password" /> </td>
</tr>
<tr>
<td></td>
<td><input type="submit" value="提交" /> </td>
<tr>
</form>
</BODY>
</HTML> `)

	case 1: // http/1.0 重定向
		// overwrite the client http version
		r.Proto = "HTTP/1.0"
		r.ProtoMajor = 1
		// r.ProtoMinor = 0

		if r.ProtoMinor == 0 {
			http.Redirect(w, r, "/login.go?mode=11", http.StatusSeeOther)
		} else {
			http.Redirect(w, r, "/login.go?mode=12", http.StatusSeeOther)
		}
	case 2: // http/1.1 重定向
		// overwrite the client http version
		r.Proto = "HTTP/1.1"
		r.ProtoMajor = 1
		r.ProtoMinor = 1

		http.Redirect(w, r, "/login.go?mode=2", http.StatusSeeOther)

	case 3: // javescript 跳转
		portalUrl := "http://" + FAKE_IP.To4().String() + "/login.go?mode=3"
		tmpl, err := template.New("foo").Parse(
			`<html>
<script language="javascript">
location.replace("");
</script>
<body>
<p align="center">
<a href=">">Welcome. If the browser cannot be redirected automatically, please click here, Thank you.</a>
</body>
</html>`)
		// type TemplateData struct {
		//   PortalUrl string
		// }
		if err != nil {
			fmt.Println("Template parse error. ", err)
			fmt.Fprint(w, "Template parse error.")
			return
		}
		err = tmpl.Execute(w, struct{ PortalUrl string }{PortalUrl: portalUrl})
		if err != nil {
			fmt.Println("Template execute error. ", err)
			fmt.Fprint(w, "Template execute error.")
		}

	case 4: // wispr
		// wispr specification 在网上已经找不到了。
		// 参考 https://github.com/wichert/wispr/blob/master/src/wispr/__init__.py 这里的代码大概猜测一下
		// https://docs.microsoft.com/en-us/windows-hardware/drivers/mobilebroadband/wispr-authentication
		wisprLoginUrl := "http://" + FAKE_IP.To4().String() + "/login.go"
		// 临时把method改一下，不然method == "GET" 的时候Redirect方法写一条链接进去网页内容里面去，我们不要写这个
		method := r.Method
		r.Method = "POST"
		http.Redirect(w, r, wisprLoginUrl, http.StatusFound)
		r.Method = method
		fmt.Fprint(w,
			`<HTML>
<!--
<?xml version="1.0" encoding="UTF-8"?>
<WISPAccessGatewayParam xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
  xsi:noNamespaceSchemaLocation="http://www.acmewisp.com/WISPAccessGatewayParam.xsd">
  <Redirect>
    <VersionHigh>1.0</VersionHigh>
    <VersionLow>0.0</VersionLow>
    <AccessProcedure>1.0</AccessProcedure>
    <AccessLocation>Hotel Guest Network</AccessLocation>
    <LocationName>Hotel Test</LocationName>
    <LoginURL>`+wisprLoginUrl+`</LoginURL>
    <MessageType>100</MessageType>
    <ResponseCode>0</ResponseCode>
  </Redirect>
</WISPAccessGatewayParam>
-->
</HTML>`

	case 5: // http head refresh 重定向   非标准，但大多浏览器都支持
		// req.Header.Add("If-None-Match", `W/"wyzzy"`)
		portalUrl := "http://" + FAKE_IP.To4().String() + "/login.go?mode=5"
		// w.Header().Set("Refresh", "0; url="+portalUrl)   // android  不支持这个，但支持下面的。IE 两个都支持
		fmt.Fprint(w, `<HTML><HEAD><TITLE>portal</TITLE><meta http-equiv="refresh" content="0; url=`+portalUrl+`"></HEAD></HTML>`)
	}
}

func PrintLocalIP() {
	ifaces, err := net.Interfaces()
	// handle err
	if err != nil {
		fmt.Println("获取不到本地ip地址\n")
		return
	}
	for _, i := range ifaces {
		addrs, err := i.Addrs()
		if err != nil {
			fmt.Println("获取不到本地ip地址\n")
			return
		}
		for _, address := range addrs {
			// check the address type and if it is not a loopback the display it
			if ipnet, ok := address.(*net.IPNet); ok && !ipnet.IP.IsLoopback() {
				if ipnet.IP.To4() != nil {
					fmt.Println(`可以通过这个地址访问 http:\\` + ipnet.IP.String())
				}
			}
		}
	}
}

//------------------------------------------------------------------------------

func main() {
	go func(msg string) {
		var h PortalServer
		// http.ListenAndServe("localhost:4000", &h)
		fmt.Println("\nhttp portal服务器已经启动")
		PrintLocalIP()
		http.ListenAndServe(":80", &h)
	}("http go coroutine")

	// ----------------------
	var servers []string
	args := flag.Args()

	if len(args) > 0 {
		servers = append(servers, args...)
	} else {
		servers = DEFAULT_SERVERS
	}

	LOG.Printf("Servers: %s", servers)

	proxyer := Proxy{servers}

	if err := dns.ListenAndServe("0.0.0.0:53", "udp", proxyer); err != nil {
		LOG.Fatal(err)
	}
}