pdqhash-go

git clone git://archive.git.mtrnord.blog/MTRNord/pdqhash-go.git
Log | Files | Refs | LICENSE

scanner.go (2769B)

      1 package main
      2 
      3 import (
      4 	"flag"
      5 	"log"
      6 	"os"
      7 	"path/filepath"
      8 	"runtime/pprof"
      9 
     10 	pdq "github.com/MTRNord/pdqhash-go"
     11 	"github.com/MTRNord/pdqhash-go/types"
     12 	"github.com/davidbyttow/govips/v2/vips"
     13 	"github.com/h2non/filetype"
     14 )
     15 
     16 // This is an example. It is not meant to be run in prod.
     17 
     18 func processFolder(filename string, detailed bool) error {
     19 	pdqhasher := pdq.NewPDQHasher()
     20 
     21 	numPDQHash := 0
     22 	var prevHash *types.Hash256
     23 	err := filepath.Walk(filename, func(fullPath string, item os.FileInfo, err error) error {
     24 		if err != nil {
     25 			return err
     26 		}
     27 		if !item.IsDir() {
     28 			// Check if file is an image
     29 			filetypeRef, err := filetype.MatchFile(fullPath)
     30 			if err != nil {
     31 				log.Fatal(err)
     32 			}
     33 			if filetypeRef.MIME.Type == "image" {
     34 				hashAndQuality := pdqhasher.FromFile(fullPath)
     35 				delta := 0
     36 				if numPDQHash == 0 {
     37 					delta = 0
     38 				} else {
     39 					delta = hashAndQuality.Hash.HammingDistance(prevHash)
     40 				}
     41 
     42 				if detailed {
     43 					log.Printf("hash=%s,norm=%d,delta=%d,quality=%d,filename=%s", hashAndQuality.Hash.String(), hashAndQuality.Hash.HammingNorm(), delta, hashAndQuality.Quality, fullPath)
     44 				} else {
     45 					log.Printf("%s,%d,%s", hashAndQuality.Hash.String(), hashAndQuality.Quality, fullPath)
     46 				}
     47 
     48 				prevHash = hashAndQuality.Hash
     49 				numPDQHash++
     50 			}
     51 		}
     52 		return nil
     53 	})
     54 	return err
     55 }
     56 
     57 func processFile(filename string, detailed bool) {
     58 	pdqhasher := pdq.NewPDQHasher()
     59 
     60 	hashAndQuality := pdqhasher.FromFile(filename)
     61 	delta := 0
     62 	if detailed {
     63 		log.Printf("hash=%s,norm=%d,delta=%d,quality=%d,filename=%s", hashAndQuality.Hash.String(), hashAndQuality.Hash.HammingNorm(), delta, hashAndQuality.Quality, filename)
     64 	} else {
     65 		log.Printf("%s,%d,%s", hashAndQuality.Hash.String(), hashAndQuality.Quality, filename)
     66 	}
     67 }
     68 
     69 func main() {
     70 	var folder string
     71 	var detailedOutput bool
     72 
     73 	flag.StringVar(&folder, "folder", "", "Folder to scan")
     74 	flag.BoolVar(&detailedOutput, "detailed", false, "Detailed output")
     75 	var cpuprofile = flag.String("cpuprofile", "", "write cpu profile to file")
     76 
     77 	flag.Parse()
     78 	if *cpuprofile != "" {
     79 		f, err := os.Create(*cpuprofile)
     80 		if err != nil {
     81 			log.Fatal(err)
     82 		}
     83 		pprof.StartCPUProfile(f)
     84 		defer pprof.StopCPUProfile()
     85 	}
     86 
     87 	// Check if folder exists and is a folder
     88 	fileInfo, err := os.Stat(folder)
     89 	if err != nil {
     90 		log.Fatal(err)
     91 	}
     92 
     93 	vips.LoggingSettings(nil, vips.LogLevelMessage)
     94 	vips.Startup(&vips.Config{
     95 		ConcurrencyLevel: 0,
     96 		MaxCacheFiles:    5,
     97 		MaxCacheMem:      50 * 1024 * 1024,
     98 		MaxCacheSize:     100,
     99 		ReportLeaks:      false,
    100 		CacheTrace:       false,
    101 		CollectStats:     false,
    102 	})
    103 	defer vips.Shutdown()
    104 	if !fileInfo.IsDir() {
    105 		processFile(folder, detailedOutput)
    106 	} else {
    107 		err := processFolder(folder, detailedOutput)
    108 		if err != nil {
    109 			log.Fatal(err)
    110 		}
    111 	}
    112 }