package blacklist

import (
	"bytes"
	"errors"
	"log"
	"net"
	"strconv"
	"strings"

	db "imuslab.com/arozos/mod/database"
)

/*

	ArozOS Blacklist Module
	Author: tobychui

	This module record the IP blacklist of users trying to enter the
	system without permission

*/

type BlackList struct {
	Enabled  bool
	database *db.Database
}

func NewBlacklistManager(sysdb *db.Database) *BlackList {
	sysdb.NewTable("ipblacklist")

	blacklistEnabled := false
	if sysdb.KeyExists("ipblacklist", "enable") {
		err := sysdb.Read("ipblacklist", "enable", &blacklistEnabled)
		if err != nil {
			log.Println("[Auth/Blacklist] Unable to load previous enable state from database. Using default.")
		}
	}

	return &BlackList{
		Enabled:  blacklistEnabled,
		database: sysdb,
	}
}

//Check if a given IP is banned
func (bl *BlackList) IsBanned(ip string) bool {
	if bl.Enabled == false {
		return false
	}
	if bl.database.KeyExists("ipblacklist", ip) {
		return true
	}

	//The ip might be inside as a range. Do a range search.
	//Need optimization, current implementation is O(N)
	for _, thisIpRange := range bl.ListBannedIpRanges() {
		if ipInRange(ip, thisIpRange) {
			return true
		}
	}
	return false
}

func (bl *BlackList) ListBannedIpRanges() []string {
	entries, err := bl.database.ListTable("ipblacklist")
	if err != nil {
		return []string{}
	}
	results := []string{}
	for _, keypairs := range entries {
		thisIpRange := keypairs[0]
		if string(thisIpRange) == "enable" || validateIpRange(string(thisIpRange)) != nil {
			//Reserved key field
			continue
		}
		results = append(results, string(thisIpRange))
	}
	return results
}

//Set the ban state of a ip or ip range
func (bl *BlackList) Ban(ipRange string) error {
	//Check if the IP range is correct
	err := validateIpRange(ipRange)
	if err != nil {
		return err
	}

	//Push it to the ban list
	ipRange = strings.TrimSpace(ipRange)
	ipRange = strings.ReplaceAll(ipRange, " ", "")
	return bl.database.Write("ipblacklist", ipRange, true)
}

//Unban an IP or IP range
func (bl *BlackList) UnBan(ipRange string) error {
	//Check if the IP range is correct
	err := validateIpRange(ipRange)
	if err != nil {
		return err
	}

	//Check if the ip range is banned
	if !bl.database.KeyExists("ipblacklist", ipRange) {
		return errors.New("invalid IP range given")
	}

	//Ip range exists, remove it from database
	return bl.database.Delete("ipblacklist", ipRange)
}

//Break an ip range text into independent ip strings
func breakdownIpRange(ipRange string) []string {
	ipRange = strings.ReplaceAll(ipRange, " ", "")
	err := validateIpRange(ipRange)
	if err != nil {
		return []string{}
	}
	if !strings.Contains(ipRange, "-") {
		//This is not an ip range but a single ip
		return []string{ipRange}
	}

	//Break down the IP range
	results := []string{}
	ips := strings.Split(ipRange, "-")

	subnet := ips[0][:strings.LastIndex(ips[0], ".")]
	startD := ips[0][strings.LastIndex(ips[0], ".")+1:]
	if err != nil {
		return []string{}
	}
	endD := ips[1][strings.LastIndex(ips[0], ".")+1:]
	if err != nil {
		return []string{}
	}

	startDInt, err := strconv.Atoi(startD)
	endDInt, err := strconv.Atoi(endD)

	currentDInt := startDInt
	for currentDInt < endDInt+1 {
		results = append(results, subnet+"."+strconv.Itoa(currentDInt))
		currentDInt++
	}

	return results
}

//Check if an given ip in the given range
func ipInRange(ip string, ipRange string) bool {
	ip = strings.TrimSpace(ip)
	ipRange = strings.ReplaceAll(ipRange, " ", "")
	if ip == ipRange {
		//For fields that the ipRange is the ip itself
		return true
	}

	//Try matching range
	if strings.Contains(ipRange, "-") {
		//Parse the source IP
		trial := net.ParseIP(ip)

		//Parse the IP range
		ips := strings.Split(ipRange, "-")
		ip1 := net.ParseIP(ips[0])
		if ip1 == nil {
			return false
		}
		ip2 := net.ParseIP(ips[1])
		if ip2 == nil {
			return false
		}
		if trial.To4() == nil {
			return false
		}
		if bytes.Compare(trial, ip1) >= 0 && bytes.Compare(trial, ip2) <= 0 {
			return true
		}
		return false

	}
	return false
}

//Check if the given IP Range string is actually an IP range
func validateIpRange(ipRange string) error {
	ipRange = strings.TrimSpace(ipRange)
	ipRange = strings.ReplaceAll(ipRange, " ", "")
	if strings.Contains(ipRange, "-") {
		//This is a range
		if strings.Count(ipRange, "-") != 1 {
			//Invalid range defination
			return errors.New("Invalid ip range defination")
		}
		ips := strings.Split(ipRange, "-")
		//Check if the starting IP and ending IP are both valid
		if net.ParseIP(ips[0]) == nil {
			return errors.New("Starting ip is invalid")
		}

		if net.ParseIP(ips[1]) == nil {
			return errors.New("Ending ip is invalid")
		}

		//Check if the ending IP is larger than the starting IP
		startingIpInt, _ := strconv.Atoi(strings.ReplaceAll(ips[0], ".", ""))
		endingIpInt, _ := strconv.Atoi(strings.ReplaceAll(ips[1], ".", ""))

		if startingIpInt >= endingIpInt {
			return errors.New("Invalid ip range: Starting IP is larger or equal to ending ip")
		}

		//Check if they are in the same subnet
		startSubnet := ips[0][:strings.LastIndex(ips[0], ".")]
		endSubnet := ips[1][:strings.LastIndex(ips[1], ".")]

		if startSubnet != endSubnet {
			//They are not in the same subnet
			return errors.New("IP range subnet mismatch")
		}

	} else {
		//This is a single IP instead of range. Check if it is a valid IP addr
		if net.ParseIP(ipRange) != nil {
			//Ok
			return nil
		} else {
			return errors.New("Invalid ip given")
		}
	}

	return nil
}