package main

import (
	"fmt"
	"net/http"
	"sync"
)

// based on: https://github.com/hakluke/hakoriginfinder

import (
	"bufio"
	"crypto/tls"
	"flag"
	"io/ioutil"
	"log"
	"os"
	"time"
)

// Make HTTP request, check response
func worker(ips <-chan string, resChan chan<- string, wg *sync.WaitGroup, client *http.Client, hostname string) {
	defer wg.Done()

	for ip := range ips {
		urls := []string{"http://" + ip, "https://" + ip}

		for _, url := range urls {
			req, err := http.NewRequest("GET", url, nil)
			if err != nil {
				fmt.Println("Error creating HTTP request", err)
				continue
			}

			req.Header.Set("Host", hostname)
			req.Header.Del("User-Agent")

			resp, err := client.Do(req)
			if err != nil {
				continue
			}

			body, err := ioutil.ReadAll(resp.Body)

			if err == nil && string(body) != "" {
				resChan <- "MATCH " + url
			} else {
				resChan <- "MAYBE " + url
			}
		}
	}
}

func main() {
	// Set up CLI flags
	workers := flag.Int("t", 32, "numbers of threads")
	hostname := flag.String("h", "", "hostname of site, e.g. www.hakluke.com")
	flag.Parse()

	// Sanity check, print usage if no hostname specified
	if *hostname == "" {
		fmt.Println("A list of IP addresses must be provided via stdin, along with a hostname of the website you are trying to find the origin of.\n\nE.g. prips 1.1.1.0/24 | hakoriginfinder -h www.hakluke.com\n\nOptions:")
		flag.PrintDefaults()
		os.Exit(2)
	}

	// IP addresses are provided via stdin
	scanner := bufio.NewScanner(os.Stdin)

	// this channel will contain the ip addresses from stdin
	ips := make(chan string)

	// this is the channel used to push a response to
	resChan := make(chan string)

	// this channel indicates when the jobs are done
	done := make(chan struct{})

	// Set up Transport (disable SSL verification)
	transport := &http.Transport{
		TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
	}

	// Set up HTTP client
	var client = &http.Client{
		Timeout:   time.Second * 10,
		Transport: transport,
	}

	// Set up waitgroup
	var wg sync.WaitGroup
	wg.Add(*workers)

	// Wait for workers to be done, then close the "done" channel
	go func() {
		wg.Wait()
		close(done)
	}()

	// Fire up workers
	for i := 0; i < *workers; i++ {
		go worker(ips, resChan, &wg, client, *hostname)
	}

	// Add ips from stdin to ips channel
	go func() {
		for scanner.Scan() {
			ips <- scanner.Text()
		}
		if err := scanner.Err(); err != nil {
			log.Println(err)
		}
		close(ips)
	}()

	// print responses from response channel, or finish
	for {
		select {
		case <-done:
			return
		case res := <-resChan:
			// print results
			fmt.Println(res)
		}
	}
}