afr0dita/afr0dita.go

/*
Description: Tool to save images.
Auth0r: sml@lacashita.com

Use it only as educational purpose.
*/


package main

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

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

type Message struct {
	Msj string
}

func main() {
	var filedb string
	var homepath = os.Getenv("HOME")
    if os.Getenv("afr0dita") == "" {
        filedb = homepath+"/afr0dita.db"
    } else {
        filedb = os.Getenv("afr0dita")
    }
 	// Create ~/.afr0ditafiles folder to save screenshots
 	if _, err := os.Stat(homepath+"/.afr0ditafiles"); os.IsNotExist(err) {
		err = os.Mkdir(homepath+"/.afr0ditafiles", 0755)
		if err != nil {
			log.Fatal(err)
		}
	}
	// Link ~/.afr0ditafiles to screenshots/ folder.
	if _, err := os.Stat("screenshots"); os.IsNotExist(err) {
		cmd := exec.Command("/usr/bin/ln", "-s", homepath+"/.afr0ditafiles", "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 afr0dita.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) {
	createafr0ditatable := `CREATE TABLE afr0dita (
	"id" INTEGER PRIMARY KEY AUTOINCREMENT,
	"date" TEXT,
	"url" TEXT,
	"tags" TEXT,
	"screenshot" TEXT
	);`
	statement, err := db.Prepare(createafr0ditatable)
	if err != nil {
		log.Fatal(err.Error())
	}
	statement.Exec()
	log.Println("[+] afr0dita 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("/", afr0ditaHandler(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:8001", nil)
}

func afr0ditaHandler(db *sql.DB) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		var repetido int
		tagtose := r.URL.Query()["tags"]
		tagtosearch := strings.Join(tagtose ," ")
		urltoa := r.URL.Query()["url"]
		urltoadd := strings.Join(urltoa, " ")
		repetido = dupURL(db,urltoadd)
		urltodel := r.URL.Query()["delete"]
		urltodelete := strings.Join(urltodel ," ")

		var homepath = os.Getenv("HOME")
		filename := getFilename(urltoadd)
		screenshotname := filename+".jpg"
		screenshotfullpath := homepath+"/.afr0ditafiles/"+filename+".jpg"							

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

		if urltoadd == "" {
			displayafr0dita(db, w,tagtosearch)
		} else {
			if repetido != 0 {
				m := Message{"URL is duplicated."}	
				printMessage(w,m)
			} else {
				err := archiveURL(urltoadd,screenshotfullpath,w)
				if err != nil {
					m := Message{"Error with the image."}
					printMessage(w,m)
				} else {
					insertUrl(db,urltoadd,tagtosearch,screenshotname)
					m := Message{"URL added successfully."}
					printMessage(w,m)
				}
			}
		}
	}
}

func displayafr0dita(db *sql.DB, w http.ResponseWriter,tagtosearch string) {
	var Arrayurlinfo []urlinfo
	var Repetidos []string
	tmpl := template.Must(template.ParseFiles("static/html/main.html"))
	p := urlinfo{}
	if tagtosearch == "" {
		row, err := db.Query("SELECT date,url,tags,screenshot FROM afr0dita")
		if err != nil {
			log.Fatal(err)
		}
		defer row.Close()
		for row.Next() {
			var date string
			var url string
			var tags string
			var screenshot string
			row.Scan( &date, &url, &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,tags,screenshot,s}
			Arrayurlinfo = append(Arrayurlinfo, p)
		}
	} 

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

	//Repetidos store now []string without tags duplicated. 
	Repetidos = removeDuplicates(Repetidos)
	// Creates the last urlinfo struct with Repetidos.
	lastp := urlinfo{"","","","",nil}
	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, tagstoinsert string, screenshotname string) {
	currentTime := time.Now()
	date := currentTime.Format("01-02-2006")
	insertURLSQL := `INSERT INTO afr0dita(date,url,tags,screenshot) VALUES (?,?,?,?)`
	statement, err := db.Prepare(insertURLSQL)
	if err != nil {
		log.Fatalln(err.Error())
	}
	_, err = statement.Exec(date, url, tagstoinsert,screenshotname)
	if err != nil {
		log.Fatalln(err.Error())
	}
}

// Delete URL.
func deleteUrl(db *sql.DB, url string) {
	deleteURLSQL := `DELETE FROM afr0dita 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 afr0dita 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 afr0dita 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 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 afr0dita 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.
func archiveURL(URL string, fileName string,w http.ResponseWriter) error {
	// Set http client to make the request with custom User-Agent.
	client := &http.Client{}
	req, _ := http.NewRequest("GET", URL, nil)
	req.Header.Set("User-Agent", "Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:106.0) Gecko/20100101 Firefox/106.0")
	// Makes request.
	response, err := client.Do(req)
	if err != nil {
		return err
	}
	defer response.Body.Close()

	// Check the URL and if the mime is an image or not. If not return an error.
	bytes, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return err
	}
	mimeType := http.DetectContentType(bytes)
	if mimeType != "image/jpg" && mimeType != "image/jpg" && mimeType != "image/png" && mimeType != "image/bmp" && mimeType != "image/jpeg"  {
		m := Message{"URL is not an image."}
		printMessage(w,m)
		err := errors.New("URL is not an image")
		return err
	} 

	// If mime is an image makes a request to get the image.
	response, err = client.Do(req)
	if err != nil {
		return err
	}
	defer response.Body.Close()

	if response.StatusCode != 200 {
		m := Message{"Returning not 200 Code"}
		printMessage(w,m)
	}
	//Create a empty file
	file, err := os.Create(fileName)
	if err != nil {
		m := Message{"Error creating the image."}
		printMessage(w,m)
		return err
	}
	defer file.Close()

	//Write the bytes to the file
	_, err = io.Copy(file, response.Body)
	if err != nil {
		m := Message{"Error copying the image."}
		printMessage(w,m)
		return err
	}
	return nil

}