//----------------------------------------------------------------------------
// Получение и обработка данных о пользовательских сессиях (авторизация)
// при соединении с почтовой системой. И блокирование нежелательных адресов.
// ---------------------------------------------------------------------------
// Версия 2 (ipset)
// ---------------------------------------------------------------------------

package main

import (
    "log"
	"strings"
	"net/netip"
    "time"
    "context"
    "fmt"
    "os"
    "io/ioutil"
    "path/filepath"
    	
    // "github.com/coreos/go-iptables/iptables"
    "github.com/gonetx/ipset"
   	"github.com/jackc/pgx/v5"
    // "github.com/jackc/pgtype"
)

// проверяем наличие ipset  системе
func CheckIPSet() {
	if err := ipset.Check(); err != nil {
		panic(err)
	}
}

type Alarm struct {
    Id int32
    Mail_account_id int32
    Remote_ip netip.Prefix
    Country_code string
    Last_time string
    Last_check_timestamp time.Time
    Action_timestamp time.Time
    Alarm_type int
    Alarm_action bool
    Hostname string
}

// PIDFile stored the process id
type PIDFile struct {
    path string
}

var PID_file_path = "/var/run/alarmaction.pid"

// Remove the pid file
// func (file PIDFile) PIDRemove() error {
    // return os.Remove(file.path)
// }

func ProcessExit(exit_code int) {
    if exit_code == 0 {
        os.Remove(PID_file_path)
    }
    os.Exit(exit_code)
}

// Выборка записей по конкретному IP
// запросы разные в зависимости от типа тревоги
func PgSelectAlarmRemoteIP(alarm Alarm) pgx.Rows {
    var query string

    conn, err := pgx.Connect(context.Background(), os.Getenv("DATABASE_URL"))
    // fmt.Println(*conn)
    if err != nil {
        fmt.Fprintf(os.Stderr, "Unable to connect to database: %v\n", err)
        ProcessExit(1)
    }

    switch alarm.Alarm_type {
        case 4:
           query = fmt.Sprintf("select id, remote_ip from alarm where alarm_type = %v and \"%v_action_timestamp\" is NULL",  alarm.Alarm_type, alarm.Hostname)
    }
    log.Println(query)
    rows, err_select := conn.Query(context.Background(), query)

    if err_select != nil {
        log.Println("Query alarms failed PgSelectAlarmRemoteIP:", err_select)
        ProcessExit(1)
    }
    if len(rows.FieldDescriptions()) == 0 {
        ProcessExit(1)
    }
    
    // log.Println("Selected an:",len(rows.FieldDescriptions()), "records")
    
    defer conn.Close(context.Background())
    return rows
}
// Проверка на соответствие адреса записи в filter
func PgSelectFilterIP(ip netip.Prefix) int32 {
    var query string
    var id int32
    conn, err := pgx.Connect(context.Background(), os.Getenv("DATABASE_URL"))
    // fmt.Println(*conn)
    if err != nil {
        fmt.Fprintf(os.Stderr, "Unable to connect to database: %v\n", err)
        ProcessExit(1)
    }

    // проверяем адрес на совпадение с фильтром (для исключения)
    // select * from filter where remote_ip >>= inet '209.85.161.20/32'
    query = fmt.Sprintf("select id from filter where remote_ip >>= inet '%v'", ip)
    log.Println(query)
    err_select := conn.QueryRow(context.Background(), query).Scan(&id)

    if err_select != nil {
    	   log.Println("Select filter for remote ip", ip, "error:", err)
    	   if err_select.Error()	 == "no rows in result set" {
           id = 0
        } else {
            ProcessExit(1)
        }
    }
    return id
}
func PgUpdateAlarm(conn *pgx.Conn, alarm Alarm, id int32) int32 {
    log.Println("Update alarm data", alarm)
    t := time.Now().Format("2006-01-02 15:04:05")
    query := fmt.Sprintf("update alarm set \"%v_action_timestamp\"='%v' where id=%v returning id", alarm.Hostname, t, id)
    log.Println("Update alarm query:", query)
    err := conn.QueryRow(context.Background(), query).Scan(&id)
    if err != nil {
	    log.Printf("Update alarm data failed: %v\n", err)
    }

    return id
}

// func IpTablesCreateChain (sChain string) bool {
    // ipt, err := iptables.New()
    // if err != nil {
        // log.Printf("Failed to new up an IPtables intance. ERROR: %v", err)
        // ProcessExit(1)
    // }
    // var chain_exists bool
    // chain_exists, err = ipt.ChainExists("filter", sChain)
    // if err != nil {
        // log.Printf("Failed to checking chain (%v). ERROR: %v", sChain	, err)
        // return false
    // }
    // if chain_exists {
        // return true
    // } else {
        // err = ipt.NewChain("filter", sChain)
        // if err != nil {
            // log.Printf("Failed to creating chain (%v). ERROR: %v", sChain	, err)
            // return false
        // }
	    // err = ipt.AppendUnique("filter", "INPUT", "-j", sChain)
	    // if err != nil {
		    // log.Printf("Failed to append chain. ERROR: %v", err)
		    // return false
	    // }
	// }
// 
    // return true
    // 
// }

// func BlockIP(ip netip.Prefix, sChain string, dChain string) bool {
	// // Some default chain names
// 
	// // Get a new iptables interface
    // ipt, err := iptables.New()
    // if err != nil {
        // log.Printf("Failed to new up an IPtables intance. ERROR: %v", err)
        // ProcessExit(1)
    // }
	// // Build out the ipstring(add /32 to the end)
	// ipstr := fmt.Sprintf("%s", ip)
    // res, _ := ipt.Exists("filter", sChain, "-s", ipstr, "-j", dChain)
    // if !res {
	    // // Use the appendUnique method to put this in iptables, but only once
	    // err = ipt.AppendUnique("filter", sChain, "-s", ipstr, "-j", dChain)
	    // if err != nil {
		    // log.Printf("Failed to ban an ip(%v). ERROR: %v", ipstr	, err)
		    // return false
	    // }
	// } else {
       	// log.Printf("Failed to ban an ip(%v). ERROR: Address already banned", ipstr)
    // }
// 
	// // Since we made it here, we won
	// return true
// }
func BlockIP(ip netip.Prefix, sChain string, dChain string) bool {
	// create test set even it's exist
	// set, err := ipset.New("auto_blocked", ipset.HashIp, ipset.Exist(true), ipset.Timeout(time.Hour))
    ip_address := ip.String()
    set, err := ipset.New("auto-blocked-ipset", ipset.HashIp, ipset.Exist(true))
	// output: test
    if err != nil {
        log.Printf("Failed to create new ipset rule. ERROR: %v", err)
        // return false
    } else {
	    log.Printf("New ipset rule %v was created", set.Name())
	}

	// _ = set.Flush()

	// _ = set.Add("1.1.1.1", ipset.Timeout(time.Hour))
	err = set.Add(ip_address)
    if err != nil {
        log.Printf("Failed to added new %v to set. ERROR: %v", ip_address, err)
        return false
    }
    
	ok, err := set.Test(ip_address)
	// output: true
	log.Println(ok)
    if err != nil {
        log.Printf("Failed to added new %v to set. ERROR: %v", ip_address, err)
        return false
    }

	// ok, _ = set.Test("1.1.1.2")
	// output: false
	// log.Println(ok)

	// info, _ := set.List()
	// output: &{test hash:ip 4 family inet hashsize 1024 maxelem 65536 timeout 3600 216 0 [1.1.1.1 timeout 3599]}
	// log.Println(info)

	// _ = set.Del("1.1.1.1")

	// _ = set.Destroy()
   	return true
}

// just suit for linux
func processExists(pid string) bool {
    if _, err := os.Stat(filepath.Join("/proc", pid)); err == nil {
        return true
    }
    return false
}
 
func checkPIDFILEAlreadyExists(path string) error {
    if pidByte, err := ioutil.ReadFile(path); err == nil {
        pid := strings.TrimSpace(string(pidByte))
        if processExists(pid) {
            return fmt.Errorf("ensure the process:%s is not running pid file:%s", pid, path)
        }
    }
    return nil
}
 
// NewPIDFile create the pid file 
// path specified under production pidfile, file content pid
func NewPIDFile(path string) (*PIDFile, error) {
    if err := checkPIDFILEAlreadyExists(path); err != nil {
        return nil, err
    }
 
    if err := os.MkdirAll(filepath.Dir(path), os.FileMode(0755)); err != nil {
        return nil, err
    }
    if err := ioutil.WriteFile(path, []byte(fmt.Sprintf("%d", os.Getpid())), 0644); err != nil {
        return nil, err
    }
    return &PIDFile{path: path}, nil
}
 
func main() {
    CheckIPSet()
    // Создадим и откроем на запись файл для логов
    var alarm_log_file string
    if os.Getenv("ALARM_LOG_FILE") == "" {
        alarm_log_file = "alarmaction.log"
    } else {
        alarm_log_file = os.Getenv("ALARM_LOG_FILE")
    }
    
    f, err := os.OpenFile(alarm_log_file, os.O_RDWR | os.O_CREATE | os.O_APPEND, 0666)
    if err != nil {
        log.Fatalf("error opening file: %v", err)
    }
    defer f.Close()

    log.SetOutput(f)

    // Создание PID  файла
    _, err = NewPIDFile(PID_file_path)
    if err != nil {
        log.Println("error to create the pid file failed:", err.Error())
        ProcessExit(1)
    }
    
    // Проверка переменных окружения
    if os.Getenv("PGHOST") == "" {
      fmt.Println("Send error: make sure environment variables `PGHOST` was set")
      ProcessExit(1)
    }
    if os.Getenv("PGUSER") == "" {
      fmt.Println("Send error: make sure environment variables `PGUSER` was set")
      ProcessExit(1)
    }
    if os.Getenv("PGPASSWORD") == "" {
      fmt.Println("Send error: make sure environment variables `PGPASSWORD` was set")
      ProcessExit(1)
    }
    if os.Getenv("PGDATABASE") == "" {
      fmt.Println("Send error: make sure environment variables `PGDATABASE` was set")
      ProcessExit(1)
    }
    // if os.Getenv("HOSTNAME") == "" {
      // fmt.Println("Send error: make sure environment variables `HOSTNAME` was set")
      // ProcessExit(1)
    // }
    
    hostname, err := os.Hostname()
	if err != nil {
		fmt.Println(err)
		ProcessExit(1)
	}
    
    conn, err := pgx.Connect(context.Background(), os.Getenv("DATABASE_URL"))
    // fmt.Println(*conn)
    if err != nil {
        fmt.Fprintf(os.Stderr, "Unable to connect to database: %v\n", err)
        ProcessExit(1)
    }
	
    sChain := "auto_blocked"
    dChain := "DROP"

    // if !IpTablesCreateChain(sChain) {
        // log.Println("Iptables chain", sChain, "does not exists")
        // ProcessExit(1)
    // }
    
    var alarm Alarm
    
    alarm.Hostname = strings.Split(hostname, ".")[0]
    alarm.Action_timestamp = time.Unix(time.Now().Unix(), 0)
    alarm.Alarm_type = 4
    
    res := PgSelectAlarmRemoteIP(alarm)
    var id int32
    var ip netip.Prefix
    var filter_id int32
    for res.Next() {
        values, _ := res.Values()
        id = values[0].(int32)
        ip = values[1].(netip.Prefix)
        filter_id = PgSelectFilterIP(ip)
        log.Println(filter_id)
        if filter_id > 0 {
            log.Println("The ip address", ip, "is conained in filter table")
            log.Println(PgUpdateAlarm(conn, alarm, id))
            // continue
        } else {
            fmt.Println(id, ip)
            if BlockIP(ip, sChain, dChain) {
                log.Println("The ip address", ip, "was blocked")            
                alarm.Alarm_action = true
                log.Println(PgUpdateAlarm(conn, alarm, id))
            }
        }
    }
    ProcessExit(0)
}