bserver/bserver.go

package main

import (
	"database/sql"
	"fmt"
	"html/template"
	"log"
	"net/http"
	"os"
	"os/exec"
	"strings"
	"time"
	_ "github.com/mattn/go-sqlite3"
	"context"
	"io/ioutil"
	"github.com/chromedp/chromedp"
	"crypto/md5"
    "encoding/hex"
)

// Start Web PART
type urlinfo struct {
	Date     string
	Url      string
	Description      string
	Tags string
	Screenshot string
	Separated []string
	Duplicated []string
}

type Message struct {
	Msj string
}

func main() {
	var filedb string
	var homepath = os.Getenv("HOME")
    if os.Getenv("babil0nia_DB") == "" {
        filedb = homepath+"/babil0nia.db"
    } else {
        filedb = os.Getenv("babil0nia_DB")
    }
 	// Create ~/.babil0niafiles folder to save screenshots
 	if _, err := os.Stat(homepath+"/.babil0niafiles"); os.IsNotExist(err) {
		err = os.Mkdir(homepath+"/.babil0niafiles", 0755)
		if err != nil {
			log.Fatal(err)
		}
	}
	// Link ~/.babil0niafiles to screenshots/ folder.
	if _, err := os.Stat("screenshots"); os.IsNotExist(err) {
		cmd := exec.Command("/usr/bin/ln", "-s", homepath+"/.babil0niafiles", "screenshots")
		_, err := cmd.Output()
		if err != nil {
			log.Fatal(err)
		}
	}
	checkDB(filedb)
    sqliteDatabase, _ := sql.Open("sqlite3", filedb)
    defer sqliteDatabase.Close()
	startServer(sqliteDatabase)
	}

// Func to check if DB exists, if not create it.
func checkDB(filedb string) {
	_, err := os.Stat(filedb)
	//If file htask.db doesnt exist.
	if err != nil {
		fmt.Println("[+]Creating Database\n")
		file, err := os.Create(filedb)
		if err != nil {
			log.Fatal(err.Error())
		}
		file.Close()
		sqlDB, _ := sql.Open("sqlite3",filedb)
		defer sqlDB.Close()
		createTable(sqlDB)
	}
}

// Func to create DB.
func createTable(db *sql.DB) {
	createbabil0niatable := `CREATE TABLE babil0nia (
	"id" INTEGER PRIMARY KEY AUTOINCREMENT,
	"date" TEXT,
	"url" TEXT,
	"description" TEXT,
	"tags" TEXT,
	"screenshot" TEXT
	);`
	statement, err := db.Prepare(createbabil0niatable)
	if err != nil {
		log.Fatal(err.Error())
	}
	statement.Exec()
	log.Println("[+] babil0nia table created")
}

func startServer(db *sql.DB) {
	// All requests to /screenshots/ will serve the static content of screenshots/
	fs := http.FileServer(http.Dir("./screenshots"))
   	http.Handle("/screenshots/", http.StripPrefix("/screenshots/", fs))

    http.HandleFunc("/", babil0niaHandler(db))
	// All requests to /styles/css will serve static/css/bootstrap.min.css.
	http.HandleFunc("/styles/css", func(response http.ResponseWriter, request *http.Request) {
		http.ServeFile(response, request, "static/css/bootstrap.min.css")
	})

	http.HandleFunc("/add", func(response http.ResponseWriter, request *http.Request) {
		http.ServeFile(response, request, "static/html/add.html")
	})

	http.ListenAndServe("127.0.0.1:8002", nil)
}

func babil0niaHandler(db *sql.DB) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		enabled := len(os.Args)
		var repetido int
		tagtose := r.URL.Query()["tags"]
		tagtosearch := strings.Join(tagtose ," ")
		wordtose := r.URL.Query()["wordtosearch"]
		wordtosearch := strings.Join(wordtose," ")
		urltoa := r.URL.Query()["url"]
		urltoadd := strings.Join(urltoa, " ")
		descript := r.URL.Query()["description"]
		descriptoinsert := strings.Join(descript," ")
		repetido = dupURL(db,urltoadd)
		urltodel := r.URL.Query()["delete"]
		urltodelete := strings.Join(urltodel ," ")
		urltoar := r.URL.Query()["archive"]
		urltoarchive := strings.Join(urltoar ," ")
		existeurl := dupURL(db,urltoarchive)
		existdownload := checkDownload(db,urltoarchive)

		if urltoarchive != "" && enabled > 1 && os.Args[1] == "-a" && existeurl != 0 && existdownload == ""{ 
			var homepath = os.Getenv("HOME")
			filename := getFilename(urltoarchive)
			screenshotname := filename+".jpg"
			screenshotfullpath := homepath+"/.babil0niafiles/"+filename+".jpg"
							
			err, filename := archiveURL(db,urltoarchive,screenshotfullpath,w)
			if err != nil {
				fmt.Printf("ERROR UPDATING")
			} else {
				updateUrl(db,screenshotname,urltoarchive)
			}	
		} 

		if urltoarchive != "" && enabled < 2 || urltoarchive != "" && enabled > 1 && os.Args[1] != "-a"  {
			m := Message{"Archive is disabled."}
			printMessage(w,m)
		}

		if urltodelete != "" {
			imageURLname := getImageName(db,urltodelete)
			filetodelete := "screenshots/"+imageURLname
			os.Remove(filetodelete)
			deleteUrl(db,urltodelete)
		}

		if urltoadd == "" && descriptoinsert == "" {
			displaybabil0nia(db, w,tagtosearch,wordtosearch)
		} else {
			if repetido != 0 {
				m := Message{"URL is duplicated."}	
				printMessage(w,m)
			} else {
				insertUrl(db,urltoadd,descriptoinsert,tagtosearch)
				m := Message{"URL added successfully."}
				printMessage(w,m)
			}
		}
	}
}

func displaybabil0nia(db *sql.DB, w http.ResponseWriter,tagtosearch string, wordtosearch string) {
	var Arrayurlinfo []urlinfo
	var Repetidos []string
	tmpl := template.Must(template.ParseFiles("static/html/main.html"))
	p := urlinfo{}
	if tagtosearch == "" && wordtosearch == "" {
		row, err := db.Query("SELECT date,url,description,tags,screenshot FROM babil0nia")
		if err != nil {
			log.Fatal(err)
		}
		defer row.Close()
		for row.Next() {
			var date string
			var url string
			var description string
			var tags string
			var screenshot string
			row.Scan( &date, &url, &description, &tags, &screenshot)
			// s is []string where tags are separated by ,
			s := strings.Split(tags, ",")
			//All s elementes are appended to Repetidos []string
			for _,element := range s {
				Repetidos = append(Repetidos,element)
			}
			p = urlinfo{date,url,description,tags,screenshot,s,nil}
			Arrayurlinfo = append(Arrayurlinfo, p)
		}
	} 

	if tagtosearch != "" && wordtosearch == "" {
		row, err := db.Query("SELECT date,url,description,tags,screenshot FROM babil0nia where tags like ?","%"+tagtosearch+"%")
		if err != nil {
			log.Fatal(err)
		}
		defer row.Close()
		for row.Next() {
			var date string
			var url string
			var description string
			var tags string
			var screenshot string
			row.Scan( &date, &url, &description, &tags, &screenshot)
			var s []string
			s = strings.Split(tags, ",")
			for _,element := range s {
				Repetidos = append(Repetidos,element)
			}
			p = urlinfo{date,url,description,tags,screenshot,s,nil}
			Arrayurlinfo = append(Arrayurlinfo, p)
		}
	}

	if tagtosearch == "" && wordtosearch != "" {
		row, err := db.Query("SELECT date,url,description,tags,screenshot FROM babil0nia where url like ? or description like ?","%"+wordtosearch+"%","%"+wordtosearch+"%")
		if err != nil {
			log.Fatal(err)
		}
		defer row.Close()
		for row.Next() {
			var date string
			var url string
			var description string
			var tags string
			var screenshot string
			row.Scan( &date, &url, &description, &tags, &screenshot)
			var s []string
			s = strings.Split(tags, ",")
			for _,element := range s {
				Repetidos = append(Repetidos,element)
			}
			p = urlinfo{date,url,description,tags,screenshot,s,nil}
			Arrayurlinfo = append(Arrayurlinfo, p)
		}
	}

	//Repetidos store now []string without tags duplicated. 
	Repetidos = removeDuplicates(Repetidos)
	// Creates the last urlinfo struct with Repetidos.
	lastp := urlinfo{"","","","","",nil,Repetidos}
	Arrayurlinfo = append(Arrayurlinfo,lastp)
	tmpl.Execute(w,Arrayurlinfo)
}

// Func to Print Messages through Web.
func printMessage(w http.ResponseWriter, msg Message) {
	tmpl := template.Must(template.ParseFiles("static/html/message.html"))
	tmpl.Execute(w, msg)
}

// Inserting URL into SQL.
func insertUrl(db *sql.DB, url string, descriptoinsert string, tagstoinsert string) {
	currentTime := time.Now()
	date := currentTime.Format("01-02-2006")
	insertURLSQL := `INSERT INTO babil0nia(date,url,description,tags) VALUES (?,?,?,?)`
	statement, err := db.Prepare(insertURLSQL)
	if err != nil {
		log.Fatalln(err.Error())
	}
	_, err = statement.Exec(date, url,descriptoinsert, tagstoinsert)
	if err != nil {
		log.Fatalln(err.Error())
	}
}

// Delete URL.
func deleteUrl(db *sql.DB, url string) {
	deleteURLSQL := `DELETE FROM babil0nia WHERE url = ?`
	statement, err := db.Prepare(deleteURLSQL)
	if err != nil {
		log.Fatalln(err.Error())
	}
	_, err = statement.Exec(url)
	if err != nil {
		log.Fatalln(err.Error())
	}
}

// Func to update an URL to add the filename of the screenshot.
func updateUrl(db *sql.DB, filename string, url string) {
	updateURLSQL := `UPDATE babil0nia set screenshot = ? WHERE url = ?`
	statement, err := db.Prepare(updateURLSQL)
	if err != nil {
		log.Fatalln(err.Error())
	}
	_, err = statement.Exec(filename,url)
	if err != nil {
		log.Fatalln(err.Error())
	}
}

// Func to check if URL is duplicated, it returns an int with the number of times that URL is repeated.
func dupURL(db *sql.DB, url string) int{
	var total int
	row, err := db.Query("SELECT count(*) as total from babil0nia where url = ?",url)
	if err != nil {
		log.Fatal(err)
	}
	defer row.Close()
	for row.Next() {
		var count int
		row.Scan( &count)
		total = count	
	}
	return total
}

// Func to check if URL has the file downloaded.
func checkDownload(db *sql.DB, url string) string{
	var total string
	row, err := db.Query("SELECT screenshot from babil0nia where url = ?",url)
	if err != nil {
		log.Fatal(err)
	}
	defer row.Close()
	for row.Next() {
		var screenshot string
		row.Scan( &screenshot)
		total = screenshot	
	}
	return total
}

// Func to get the "id" that will be used as imagename from an given URL, it returns the id.
func getImageName(db *sql.DB, url string) string{
	var imagename string
	row, err := db.Query("SELECT screenshot from babil0nia where url = ?",url)
	if err != nil {
		log.Fatal(err)
	}
	defer row.Close()
	for row.Next() {
		var screenshot string
		row.Scan( &screenshot)
		imagename = screenshot	
	}
	return imagename
}

// Func to remove duplicated strings from an slice.
func removeDuplicates(strSlice []string) []string {
    allKeys := make(map[string]bool)
    list := []string{}
    for _, item := range strSlice {
        if _, value := allKeys[item]; !value {
            allKeys[item] = true
            list = append(list, item)
        }
    }
    return list
}

// Func to create a filename using the md5 hash of the url name.
func getFilename(text string) string {
    hash := md5.Sum([]byte(text))
    return hex.EncodeToString(hash[:])
}

// Func to take one screenshot from the given URL and save it to a file. It returns error and the name of the saved (screenshot) file.
func archiveURL(db *sql.DB,urltoarchive string, filename string,w http.ResponseWriter) (error,string){
	var fail error
	ctx, cancel := chromedp.NewContext(
		context.Background(),
	)
	defer cancel()
	var buf []byte
	if err := chromedp.Run(ctx, fullScreenshot(urltoarchive, 100, &buf)); err != nil {
		m := Message{"Error downloading."}
		printMessage(w,m)
		fail = err
	}
	if fail != nil {
		fmt.Println("Fallo la descarga de la imagen asi que nos e guarda")
	} else {
		if err := ioutil.WriteFile(filename, buf, 0o644); err != nil {
		m := Message{"Error saving the image."}
		printMessage(w,m)
		fail = err
		}
	}
	return fail, filename
}

func fullScreenshot(urlstr string, quality int, res *[]byte) chromedp.Tasks {
	return chromedp.Tasks{
		chromedp.Navigate(urlstr),
		chromedp.FullScreenshot(res, quality),
	}
}