shithub: mkfaces

ref: d3e11eee5dc699fd2edc050f089eec9742931eeb
dir: mkfaces/main.go

View raw version
package main

import (
	"bytes"
	"crypto/md5"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"os/exec"
	"regexp"
	"sort"
	"strings"
)

var (
	httpClient = &http.Client{}
	base       = "/mail/fs"
)

func main() {
	home, err := os.UserHomeDir()
	if err != nil {
		log.Fatal(err)
	}
	faceBase := home + "/lib/face"
	outBase := faceBase + "/48x48x8"
	notFoundPath := faceBase + "/.notfound"
	ignorePath := faceBase + "/.ignorelist"
	dictPath := outBase + "/.dict"

	froms := map[string][md5.Size]byte{}

	fs, err := ioutil.ReadDir(base)
	if err != nil {
		log.Fatal(err)
	}
	for _, fsi := range fs {
		if !fsi.IsDir() {
			continue
		}

		msgsBase := base + "/" + fsi.Name()
		msgs, err := ioutil.ReadDir(msgsBase)
		if err != nil {
			log.Fatal(err)
		}

		for _, mi := range msgs {
			if !mi.IsDir() {
				continue
			}
			if from, err := ioutil.ReadFile(msgsBase + "/" + mi.Name() + "/from"); err != nil {
				log.Fatal(err)
			} else {
				f := strings.ToLower(string(from))
				froms[f] = md5.Sum([]byte(f))
			}
		}
	}

	if err = os.MkdirAll(outBase, 0700); err != nil {
		log.Fatal(err)
	}

	var ignoreList []*regexp.Regexp
	if s, err := ioutil.ReadFile(ignorePath); err == nil {
		for _, v := range strings.Split(string(s), "\n") {
			if v != "" {
				if r, err := regexp.Compile(v); err != nil {
					log.Fatalf("%s: %s", ignorePath, err)
				} else {
					ignoreList = append(ignoreList, r)
				}
			}
		}
	}

	notFound := make(map[string]struct{})
	if s, err := ioutil.ReadFile(notFoundPath); err == nil {
		for _, v := range strings.Split(string(s), "\n") {
			if v != "" {
				notFound[v] = struct{}{}
			}
		}
	}

	dict := make(map[string]string)
	if s, err := ioutil.ReadFile(dictPath); err == nil {
		for _, v := range strings.Split(string(s), "\n") {
			if v != "" {
				parts := strings.Split(v, " ")
				dict[parts[0]] = parts[1]
			}
		}
	}

	numTotal := len(froms)
	i := 0
	failed := 0
	saved := 0
	ignored := 0
	progress := ""

	for f, h := range froms {
		hash := fmt.Sprintf("%x", h)
		imagePath := outBase + "/" + hash

		i++
		clear := strings.Repeat("\x08", len(progress))
		progress = fmt.Sprintf("%d/%d", i, numTotal)
		fmt.Printf("%s%s", clear, progress)

		var machineUser string
		if parts := strings.Split(f, "@"); len(parts) != 2 {
			failed++
			continue
		} else {
			userParts := strings.Split(parts[0], "+")
			machineUser = fmt.Sprintf("%s/%s", parts[1], userParts[0])

			skip := false
			for _, ignore := range ignoreList {
				if ignore.MatchString(machineUser) {
					ignored++
					os.Remove(imagePath)
					delete(dict, machineUser)
					delete(notFound, machineUser)
					skip = true
				}
			}
			if skip {
				continue
			}
			if _, ok := dict[machineUser]; ok {
				continue
			}
			if _, ok := notFound[machineUser]; ok {
				continue
			}
		}

		url := fmt.Sprintf("http://gravatar.com/avatar/%s.jpg?s=48&d=404", hash)
		if res, err := httpClient.Get(url); err != nil {
			log.Fatal(err)
		} else if res.StatusCode != http.StatusOK {
			if res.StatusCode == http.StatusNotFound {
				notFound[machineUser] = struct{}{}
			}
			res.Body.Close()
			continue
		} else {
			b := new(bytes.Buffer)
			b.ReadFrom(res.Body)
			res.Body.Close()

			data := new(bytes.Buffer)

			cmd := exec.Command("/bin/jpg", "-c")
			cmd.Stdin = bytes.NewReader(b.Bytes())
			cmd.Stdout = data
			if err = cmd.Run(); err != nil {
				data.Reset()
				cmd = exec.Command("/bin/png", "-c")
				cmd.Stdin = bytes.NewReader(b.Bytes())
				cmd.Stdout = data
				err = cmd.Run()
			}

			if err != nil {
				failed++
			} else if err = ioutil.WriteFile(imagePath, data.Bytes(), 0644); err != nil {
				log.Fatal(err)
			} else {
				dict[machineUser] = hash
				saved++
			}
		}
	}

	if f, err := os.Create(notFoundPath); err != nil {
		log.Fatal(err)
	} else {
		var sorted []string
		for machineUser := range notFound {
			sorted = append(sorted, machineUser)
		}
		sort.Strings(sorted)
		for _, s := range sorted {
			fmt.Fprintf(f, "%s\n", s)
		}
		f.Close()
	}

	if f, err := os.Create(dictPath); err != nil {
		log.Fatal(err)
	} else {
		var sorted []string
		for machineUser := range dict {
			sorted = append(sorted, machineUser)
		}
		sort.Strings(sorted)
		for _, s := range sorted {
			fmt.Fprintf(f, "%s %s\n", s, dict[s])
		}
		f.Close()
	}

	fmt.Printf("%s", strings.Repeat("\x08", len(progress)))
	fmt.Printf("%d addresses\n", numTotal)
	fmt.Printf("%d faces added\n", saved)
	fmt.Printf("%d failed to decode\n", failed)
	fmt.Printf("%d ignored\n", ignored)
}