package main

import (
	"encoding/json"
	"fmt"
	"github.com/go-git/go-git/v5"
	"github.com/go-git/go-git/v5/config"
	"github.com/go-git/go-git/v5/plumbing"
	"github.com/go-git/go-git/v5/plumbing/object"
	"github.com/go-git/go-git/v5/plumbing/transport"
	"github.com/go-git/go-git/v5/plumbing/transport/http"
	_ "github.com/joho/godotenv/autoload"
	log "github.com/sirupsen/logrus"
	"gopkg.in/yaml.v2"
	"io"
	"io/ioutil"
	"net"
	"os"
	"path"
	"strings"
	"sync/atomic"
	"time"
)

var (
	blacklistQueue []net.IP
	isUpdated      = atomic.Bool{}
)

func main() {
	log.SetLevel(log.TraceLevel)
	log.Info("Starting the application")

	go updateRunner()

	// Start the application
	listens := os.Getenv("LISTEN")
	ports := strings.Split(listens, ",")
	for _, port := range ports {
		go listen(port)
	}

	// Wait forever
	select {}
}

func updateRunner() {
	// Function to update the blacklist
	// Blacklist is a file stored in git repo
	repoUrl := os.Getenv("GIT_REPO")
	user := os.Getenv("GIT_USER")
	pass := os.Getenv("GIT_PASS")
	branch := os.Getenv("GIT_BRANCH")
	fileName := os.Getenv("REPO_FILE")

	log.Tracef("Initializing the repository %s", repoUrl)

	// clone repository to local
	var auth transport.AuthMethod
	if user != "" && pass != "" {
		auth = &http.BasicAuth{
			Username: user,
			Password: pass,
		}
	}

	// delete if exists
	_ = os.RemoveAll("repo")

	repository, err := git.PlainClone("repo", false, &git.CloneOptions{
		URL:      repoUrl,
		Progress: os.Stdout,
		Auth:     auth,
	})

	if err != nil {
		panic(err)
	}

	for {
		log.Info("Checking for updates")
		if isUpdated.Load() {
			isUpdated.Store(false)
			update(repository, branch, auth, fileName)
		}

		// Sleep for 1 minute
		time.Sleep(5 * time.Second)
	}
}

func update(repository *git.Repository, branch string, auth transport.AuthMethod, fileName string) {
	type structuredBlacklist []string

	log.Info("Updating the blacklist")
	log.Debugf("Branch: %s, File: %s", branch, fileName)

	// pull
	log.Debug("Pulling the repository")
	w, err := repository.Worktree()
	if err != nil {
		panic(err)
	}

	err = w.Pull(&git.PullOptions{
		ReferenceName: plumbing.NewBranchReferenceName(branch),
		SingleBranch:  true,
		RemoteName:    "origin",
		Auth:          auth,
	})
	if err != nil && err != git.NoErrAlreadyUpToDate {
		panic(err)
	}

	// read file
	log.Debug("Reading the file")
	var fileContents []byte
	rev, err := repository.ResolveRevision(plumbing.Revision(branch))
	if err != nil {
		panic(err)
	}

	commit, err := repository.CommitObject(*rev)
	if err != nil {
		panic(err)
	}

	tree, err := commit.Tree()
	if err != nil {
		panic(err)
	}

	file, err := tree.File(fileName)
	if err != nil {
		panic(err)
	}

	fileContentsReader, err := file.Reader()
	if err != nil {
		panic(err)
	}
	defer fileContentsReader.Close()

	fileContents, err = io.ReadAll(fileContentsReader)
	if err != nil {
		panic(err)
	}

	alreadyBlacklisted := make([]string, 0)

	// update blacklist
	// parse content based on extension
	if strings.HasSuffix(fileName, ".txt") {
		alreadyBlacklisted = strings.Split(string(fileContents), "\n")
		// remove # comments
		for i, v := range alreadyBlacklisted {
			if strings.HasPrefix(v, "#") || strings.Trim(v, " #") == "" {
				alreadyBlacklisted = append(alreadyBlacklisted[:i], alreadyBlacklisted[i+1:]...)
			}
		}
	} else if strings.HasSuffix(fileName, ".json") {
		// parse json
		blacklistedStruct := structuredBlacklist{}
		err = json.Unmarshal(fileContents, &blacklistedStruct)
		if err != nil {
			panic(err)
		}
		alreadyBlacklisted = blacklistedStruct
	} else if strings.HasSuffix(fileName, ".yaml") || strings.HasSuffix(fileName, ".yml") {
		// parse yaml
		blacklistedStruct := structuredBlacklist{}
		err = yaml.Unmarshal(fileContents, &blacklistedStruct)
		if err != nil {
			panic(err)
		}
		alreadyBlacklisted = blacklistedStruct
	} else {
		panic("Unsupported file extension")
	}

	toAdd := blacklistQueue
	blacklistQueue = []net.IP{}

	isAlreadyPresent := func(s net.IP) bool {
		// check if the address is already present in the list
		for _, candidate := range alreadyBlacklisted {
			if s.String() == candidate {
				return true
			}

			if strings.Contains(candidate, "/") {
				// Parse CIDR
				_, cidr, err := net.ParseCIDR(candidate)
				if err != nil {
					panic(err)
				}
				if cidr.Contains(s) {
					return true
				}
				continue
			}

			// if ipv6
			if strings.Contains(candidate, ":") {
				// Simply compare the IP is enough for IPv6 blocking
				// Can be improved by using net.ParseIP and net.IP.Equal
				continue
			}

			// has wildcard
			if strings.Contains(candidate, "*") {
				// Check if the IP is in the range
				parts := strings.Split(candidate, ".")
				ipParts := strings.Split(s.String(), ".")
				isMatch := true
				for i, part := range parts {
					if part == "*" {
						continue
					}
					if part != ipParts[i] {
						isMatch = false
						break
					}
				}
				if isMatch {
					return true
				}
			}

			continue
		}
		return false
	}

	added := make([]net.IP, 0)
	for _, v := range toAdd {
		if !isAlreadyPresent(v) {
			alreadyBlacklisted = append(alreadyBlacklisted, v.String())
			added = append(added, v)
		}
	}

	if len(added) < 1 {
		log.Info("No new addresses to add")
		return
	}

	logMessages := strings.Builder{}
	log.Infof("Added %d new addresses", len(added))
	commitMessage := strings.Builder{}
	commitMessage.WriteString(fmt.Sprintf("Add %d new addresses\n-------------------------------\n", len(added)))
	for _, v := range added {
		logMessages.WriteString(fmt.Sprintf("%s ", time.Now().Format(time.RFC3339)))
		logMessages.WriteString(fmt.Sprintf("addr='%s'", v.String()))

		commitMessage.WriteString(v.String())
		reverseDNSNames, err := net.LookupAddr(v.String())
		if err == nil {
			line := fmt.Sprintf(" hostname='%s'", reverseDNSNames[0])
			commitMessage.WriteString(line)
			logMessages.WriteString(line)
		} else {
			log.Warnf("Failed to lookup reverse DNS for %s: %s", v.String(), err.Error())
		}

		geo, err := geoLookup(v)
		if err == nil {
			line := fmt.Sprintf(
				" country='%s' city='%s' postal='%s'",
				geo.Country.Names["en"],
				geo.City.Names["en"],
				geo.Postal.Code,
			)
			commitMessage.WriteString(line)
			logMessages.WriteString(line)
		} else {
			log.Warnf("Failed to lookup geo for %s: %s", v.String(), err.Error())
		}

		asn, err := asnLookup(v)
		if err == nil {
			line := fmt.Sprintf(
				" asn.id='%d' asn.name='%s'",
				asn.AutonomousSystemNumber,
				asn.AutonomousSystemOrganization,
			)
			commitMessage.WriteString(line)
			logMessages.WriteString(line)
		} else {
			log.Warnf("Failed to lookup ASN for %s: %s", v.String(), err.Error())
		}

		commitMessage.WriteString("\n")
		logMessages.WriteString("\n")
	}

	// get content
	var newContent []byte
	if strings.HasSuffix(fileName, ".txt") {
		newContent = []byte(strings.Join(alreadyBlacklisted, "\n"))
	} else if strings.HasSuffix(fileName, ".json") {
		newContent, err = json.Marshal(alreadyBlacklisted)
		if err != nil {
			panic(err)
		}
	} else if strings.HasSuffix(fileName, ".yaml") || strings.HasSuffix(fileName, ".yml") {
		newContent, err = yaml.Marshal(alreadyBlacklisted)
		if err != nil {
			panic(err)
		}
	} else {
		panic("Unsupported file extension")
	}

	// write to file
	log.Debug("Writing to file")
	//repository.Storer.SetConfig("core.autocrlf", "false")
	err = ioutil.WriteFile(fmt.Sprintf("repo/%s", fileName), newContent, 0644)
	if err != nil {
		panic(err)
	}

	logDir := os.Getenv("REPO_LOG_DIR")
	if logDir != "" {
		// path
		logPath := fmt.Sprintf("repo/%s/%d/%d/%d.txt", logDir, time.Now().Year(), time.Now().Month(), time.Now().Day())
		_ = os.MkdirAll(path.Dir(logPath), 0755)

		// append to log
		logFile, err := os.OpenFile(logPath, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
		if err != nil {
			_ = logFile.Close()
			panic(err)
		}

		_, err = logFile.WriteString(logMessages.String())
		if err != nil {
			_ = logFile.Close()
			panic(err)
		}

		_ = logFile.Close()
		_, err = w.Add(strings.ReplaceAll(logPath, "repo/", ""))
		if err != nil {
			panic(err)
		}
	}

	// commit
	log.Debug("Committing the changes")
	_, err = w.Add(fileName)
	if err != nil {
		panic(err)
	}

	_, err = w.Commit(commitMessage.String(), &git.CommitOptions{
		Committer: &object.Signature{
			Name:  "Blacklist Robot",
			Email: "robot+blacklist@mrlokop.dev",
			When:  time.Now(),
		},
		Author: &object.Signature{
			Name:  "Blacklist Robot",
			Email: "robot+blacklist@mrlokop.dev",
			When:  time.Now(),
		},
		All: true,
	})

	if err != nil {
		panic(err)
	}

	// push
	log.Debug("Pushing the changes")
	err = repository.Push(&git.PushOptions{
		Auth:       auth,
		Progress:   os.Stdout,
		RemoteName: "origin",
		RefSpecs: []config.RefSpec{
			config.RefSpec(fmt.Sprintf("refs/heads/%s:refs/heads/%s", branch, branch)),
		},
	})

	if err != nil {
		panic(err)
	}

	log.Info("Update complete")
}

func listen(addr string) {
	log.Infof("Listening on addr %s", addr)
	ln, err := net.Listen("tcp", addr)
	if err != nil {
		log.Fatal(err)
	}
	for {
		conn, err := ln.Accept()
		if err != nil {
			log.Fatal(err)
		}

		remoteAddr := conn.RemoteAddr()
		if isPrivateIP(remoteAddr) {
			log.Warnf("Got request from private IP %s, ignoring request", remoteAddr)
			_ = conn.Close()
			continue
		}

		// check is ip
		ip := net.ParseIP(remoteAddr.String())
		if ip == nil {
			log.Warnf("Got request from non-IP %s, ignoring request", remoteAddr)
			_ = conn.Close()
			continue
		}

		_ = conn.Close()
		log.Infof("Got request from %s, blacklisting", remoteAddr)
		blacklistQueue = append(blacklistQueue, ip)
		isUpdated.Store(true)
	}
}