package main

import (
	"context"
	"crypto/tls"
	"encoding/base64"
	"fmt"
	"golang.org/x/crypto/acme/autocert"
	"log"
	"net"
	"net/http"
	"net/http/httputil"
	"net/mail"
	"os"
	"strconv"
	"strings"
	"time"
)

var (
	targetHost          = os.Getenv("SDF_PROXY_TARGET_HOST")
	targetScheme        = os.Getenv("SDF_PROXY_TARGET_SCHEME")
	myIp                string
	myPtrAddr           string
	allowedHosts        = strings.Split(os.Getenv("SDF_PROXY_ALLOWED_HOSTS"), ",")
	tlsLetsEncryptEmail = os.Getenv("SDF_PROXY_LETS_ENCRYPT_EMAIL")
	tlsCertManager      = autocert.Manager{
		Prompt:     autocert.AcceptTOS,
		Cache:      autocert.DirCache("certs"),
		Email:      tlsLetsEncryptEmail,
		HostPolicy: autocert.HostWhitelist(allowedHosts...),
	}
)

func init() {
	setMyIpAndPtrAddr(&myIp, &myPtrAddr)
	if len(myPtrAddr) > 0 {
		tlsCertManager.HostPolicy = autocert.HostWhitelist(append(allowedHosts, myPtrAddr)...)
	}
}

// setMyIpAndPtrAddr attempts to set myIp and myPtrAddr
func setMyIpAndPtrAddr(ip, addr *string) {
	resolver := &net.Resolver{
		PreferGo: true,
		Dial: func(ctx context.Context, network, address string) (net.Conn, error) {
			d := net.Dialer{
				Timeout: time.Duration(10) * time.Second,
			}
			return d.DialContext(ctx, network, "ns1.google.com:53")
		},
	}
	ips, errTxtLookup := resolver.LookupTXT(context.Background(), "o-o.myaddr.l.google.com")
	if errTxtLookup != nil {
		log.Printf("%v\n", errTxtLookup)
		return
	}
	if len(ips) > 0 {
		*ip = ips[len(ips)-1]
		addrs, errAddrLookup := net.LookupAddr(*ip)
		if errAddrLookup != nil {
			log.Printf("%v\n", errAddrLookup)
			return
		}
		if len(addrs) > 0 {
			*addr = strings.TrimSuffix(addrs[len(addrs)-1], ".")
		}
	}
}

type httpReverseProxy struct {
	http.Server
	serveTls bool
}

func newHttpReverseProxy(addr string, serveTls bool) *httpReverseProxy {
	server := new(httpReverseProxy)
	server.Addr = addr
	server.serveTls = serveTls
	if server.serveTls {
		server.TLSConfig = &tls.Config{
			GetCertificate: tlsCertManager.GetCertificate,
			MinVersion:     tls.VersionTLS12,
		}
	}

	proxy := new(httputil.ReverseProxy)
	proxy.Transport = &http.Transport{}
	proxy.Director = func(req *http.Request) {
		req.Header.Set("X-Forwarded-For", req.RemoteAddr)
		req.Header.Set("X-Forwarded-Host", req.Host)
		req.Host = targetHost
		req.URL.Host = targetHost
		req.URL.Scheme = targetScheme
		fmt.Print("Handling request: ")
		fmt.Println(req.Header)
	}

	mux := http.NewServeMux()
	// Hostnames for which the reverse proxy should direct requests to the target host
	for _, host := range allowedHosts {
		mux.Handle(host+"/", proxy)
		mux.Handle(host+"/.well-known/acme-challenge/", tlsCertManager.HTTPHandler(nil))
	}
	// First hostname in the PTR record of the IP address which is being used by the server to connect to the internet
	// This is an easter egg and will only work if the server is also reachable by the IP address from the internet
	if len(myPtrAddr) > 0 {
		mux.HandleFunc(myPtrAddr+"/", func(w http.ResponseWriter, r *http.Request) {
			w.WriteHeader(http.StatusOK)
			w.Header().Set("Content-Type", "text/html")
			content, _ := base64.StdEncoding.DecodeString("PCFkb2N0eXBlIGh0bWw+CjxodG1sIGxhbmc9ZW4+CjxoZWFkPgo8bWV0YSBjaGFyc2V0PXV0Zi04Pgo8dGl0bGU+8J+Qh/CfpZo8L3RpdGxlPgo8c3R5bGU+CmJvZHkge2JhY2tncm91bmQtY29sb3I6ICM3MDc0NzA7fQpodG1sIHtjdXJzb3I6IHVybChkYXRhOmltYWdlL3BuZztiYXNlNjQsaVZCT1J3MEtHZ29BQUFBTlNVaEVVZ0FBQUJBQUFBQVFDQVlBQUFBZjgvOWhBQUFBQVhOU1IwSUFyczRjNlFBQUFBWmlTMGRFQVA4QS93RC9vTDJua3dBQUFBbHdTRmx6QUFBTEV3QUFDeE1CQUpxY0dBQUFBQWQwU1UxRkIrQUtIUUFKRDNNdWdCa0FBQUI1U1VSQlZEakxwVk5CRHNBZ0NDdUdRLy8vV2c0bTdLS0dtRzBxa25DQlVBb0ZzV3FPQzFNQW9ESlZiTlZRc3AydEdxZ1V6UlEyMWpKR09PMGFZd1dYZGdSQUpXYlZqaG5NSUtrUklraDZCeDNrK2c3S1c2SnJ2U1ZubThVQmVGeE9qRWZmVm9GS21abThIZExvOVBXVnF6eXNtcTllK2kvL0FQUW9aZmJneUlxQkFBQUFBRWxGVGtTdVFtQ0MpLCBhdXRvO30KPC9zdHlsZT4K")
			if _, err := w.Write(content); err != nil {
				fmt.Println(err)
				return
			}
		})
		mux.Handle(myPtrAddr+"/.well-known/acme-challenge/", tlsCertManager.HTTPHandler(nil))
	}
	// All other hostnames, which are being ignored
	mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusGatewayTimeout)
	})

	server.Handler = mux

	return server
}

// listenAndServe will call ListenAndServe or ListenAndServeTLS on httpReverseProxy to handle requests on incoming connections
func (t *httpReverseProxy) listenAndServe() {
	var err error

	fmt.Printf("Listening on %s\n", t.Addr)

	if t.serveTls {
		err = t.ListenAndServeTLS("", "")
	} else {
		err = t.ListenAndServe()
	}
	if err != nil {
		log.Fatalf("listenAndServe, Addr: %s, serveTls: %s, error: %v", t.Addr, strconv.FormatBool(t.serveTls), err)
	}
}

func main() {
	if _, err := mail.ParseAddress(tlsLetsEncryptEmail); err != nil {
		log.Fatalf("Invalid Let's Encrypt account email: %v\n", err)
	}

	fmt.Print("Allowed hosts: ")
	fmt.Println(allowedHosts)
	fmt.Print("PTR record (if any): ")
	fmt.Println(myPtrAddr)

	httpProxy := newHttpReverseProxy(":http", false)
	httpsProxy := newHttpReverseProxy(":https", true)

	go httpProxy.listenAndServe()
	httpsProxy.listenAndServe()
}