Golang-phish-update/dbupdate.go

package main

import (
	"database/sql"
	"fmt"
	"os"
	"log"
	"net/http"
	"io/ioutil"
	"encoding/json"
	"time"

	"github.com/joho/godotenv"
	"github.com/cheggaaa/pb/v3"

	_ "github.com/go-sql-driver/mysql"
)

func main() {
	// Clear the screen
	fmt.Print("\033[2J")
	// Load environment variables from .env file
	err := godotenv.Load()
	if err != nil {
		log.Fatal("Error loading .env file")
	}
	//Connect to the database
	db, err := sql.Open("mysql", os.Getenv("DB_USER")+":"+os.Getenv("DB_PASS")+"@tcp("+os.Getenv("DB_HOST")+":"+os.Getenv("DB_PORT")+")/"+os.Getenv("DATABASE"))
	if err != nil {
		log.Fatal(err)
	}
	defer db.Close()

	// Verify that the connection is successful
	err = db.Ping()
	if err != nil {
		log.Fatal(err)
	}

	// Get the start time
	StartTime := time.Now()

	// Format the start time
	formattedStartTime := StartTime.Format("02-01-2006 15:04:05")

	fmt.Println("Script started at " + formattedStartTime)
	// Perform a cleanup on DB
	deleter, err := db.Exec("DELETE FROM phish;")
	if err != nil {
		log.Fatal(err)
	}
	// Get the number of affected rows
	rowsAffected, err := deleter.RowsAffected()
	if err != nil {
		log.Fatal(err)
	}
	fmt.Printf("Deletion: Number of rows affected: %d\n", rowsAffected)

	// Perform a cleanup on DB
	aiupdate, err := db.Exec("ALTER TABLE phish AUTO_INCREMENT = 0;")
	if err != nil {
		log.Fatal(err)
	}
	// Get the number of affected rows
	rowsAffected, err = aiupdate.RowsAffected()
	if err != nil {
		log.Fatal(err)
	}
	fmt.Printf("Altering: Number of rows affected: %d\n", rowsAffected)

	// Add local records first
	insert1, err := db.Exec("INSERT INTO phish (Domain, Reporter, Evidence) VALUES ('iriseperiplo.info', 'Henk', 'https://www.virustotal.com/gui/url/318ac29af0001f6678efdd94c16d3225d78369123ebbbe5f50387b208e0ac523?nocache=1');")
	insert2, err2 := db.Exec("INSERT INTO phish (Domain, Reporter, Evidence) VALUES ('storage.com', 'Henk', 'https://www.virustotal.com/gui/url/92fef9276cd4c0433171aac2fe92d140df069ff99d724c850ca1d3ccbdd0ae9f?nocache=1');")
	if err != nil {
		log.Fatal(err)
	}
	if err2 != nil {
		log.Fatal(err)
	}
	// Get the number of affected rows
	rowsAffected, err = insert1.RowsAffected()
	if err != nil {
		log.Fatal(err)
	}
	fmt.Printf("Insert: Number of rows affected: %d\n", rowsAffected)
	// Get the number of affected rows
	rowsAffected, err = insert2.RowsAffected()
	if err != nil {
		log.Fatal(err)
	}
	fmt.Printf("Insert: Number of rows affected: %d\n", rowsAffected)

	// Define name for use later
	// name := os.Getenv("USER_OR_BOTNAME")

	// Make an HTTP request to phish.sinking.yachts dbsize endpoint
	sizereq, err := http.NewRequest("GET", "https://phish.sinking.yachts/v2/dbsize", nil)
	if err != nil {
		log.Fatal(err)
	}
	sizereq.Header.Set("X-Identity", "ultimatebot db updater via Golang")
	sizeclient := &http.Client{}
	sizeresp, err := sizeclient.Do(sizereq)
	if err != nil {
		log.Fatal(err)
	}
	defer sizeresp.Body.Close()

	// Request status code of endpoint
	fmt.Println("Status code dbsize endpoint: ", sizeresp.StatusCode)
	sizebody, err := ioutil.ReadAll(sizeresp.Body)
	if err != nil {
		log.Fatal(err)
	}

	// Unmarshal request (To be sure)
	var sizedata int64
	err = json.Unmarshal(sizebody, &sizedata)
	if err != nil {
		log.Fatal(err)
	}

	maxpbsize := sizedata

	// Make an HTTP request to phish.sinking.yachts alldomains endpoint
	req, err := http.NewRequest("GET", "https://phish.sinking.yachts/v2/all", nil)
	if err != nil {
		log.Fatal(err)
	}
	req.Header.Set("X-Identity", "ultimatebot db updater via Golang")
	client := &http.Client{}
	resp, err := client.Do(req)
	if err != nil {
		log.Fatal(err)
	}
	defer resp.Body.Close()

	// Setup current time variable, which will get used later
	currentTime := time.Now()
	formattedTime := currentTime.Format("02-01-2006 15:04:05")
	fmt.Println(formattedTime + ": Initial database table setup done")

	// Request status code of endpoint
	fmt.Println("Status code: ", resp.StatusCode)
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Fatal(err)
	}
	// Unmarshal request
	var data []string
	err = json.Unmarshal(body, &data)
	if err != nil {
		log.Fatal(err)
	}

	// Initialize a variable to keep track of the total number of rows affected
	var totalRowsAffected int64

	// Start processbar with variable from earlier with value from dbsize endpoint
	bar := pb.Full.Start64(maxpbsize)

	for _, element := range data {
		// fmt.Println(element)
		// Add domains
		inserter, err := db.Exec("INSERT INTO phish (Domain) VALUES (?)", element)
		if err != nil {
			log.Fatal(err)
		}

		rowsAffected, err := inserter.RowsAffected()
		if err != nil {
			log.Fatal(err)
		}

		totalRowsAffected += rowsAffected
		bar.Increment()
		time.Sleep(time.Millisecond)
	}
	// finish the processbar
	bar.Finish()

	// Print the totalRowsAffected (Which will be the same as max processbar)
	fmt.Printf("Total number of rows affected: %d\n", totalRowsAffected)

	// Get the end time
	EndTime := time.Now()

	// Format the end time
	formattedEndTime := EndTime.Format("02-01-2006 15:04:05")

	fmt.Println("Recent changes recieved and has been added to the database. Removed domains deleted from database. Reorded domains.")

	fmt.Println("Script ended at " + formattedEndTime)

	// Calculate the time difference
	diff := EndTime.Sub(StartTime)

	// Print the time difference
	fmt.Println("Time until completion: ", diff)

	// dummy
	// fmt.Println("Data: ", string(body))
}