docstore/pkg/scan/scanner.go

// Copyright (C) 2021 Marius Schellenberger

package scan

import (
	"errors"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"time"

	"git.giftfish.de/ston1th/docstore/pkg/core"
	"git.giftfish.de/ston1th/docstore/pkg/fs"
	"git.giftfish.de/ston1th/godrop/v2"
	"github.com/go-logr/logr"
)

const (
	pdfExt = ".pdf"

	minThreads = 1
	maxThreads = 4
)

type PDF struct {
	cmd  string
	args []string
}

var defaultLangs = []string{"eng", "deu"}

type Tesseract struct {
	cmd  string
	args []string
	env  []string
}

var (
	wsr = regexp.MustCompile(`[\t\f ]+`)
	nlr = regexp.MustCompile(`\n{3,}`)
	ws  = []byte(" ")
	nl  = []byte("\n")
)

type Scanner struct {
	log     logr.Logger
	base    string
	timeout time.Duration
	t       Tesseract
	p       PDF
}

func New(log logr.Logger, c *core.Config) (s *Scanner, err error) {
	var (
		tcmd string
		pcmd string
		lang string
	)
	tcmd, err = exec.LookPath("tesseract")
	if err != nil {
		return
	}
	pcmd, err = exec.LookPath("pdftotext")
	if err != nil {
		return
	}
	err = godrop.Unveil(tcmd, "rx")
	if err != nil {
		return
	}
	err = godrop.Unveil(pcmd, "rx")
	if err != nil {
		return
	}
	llen := len(c.Langs)
	langs := c.Langs
	if langs == nil || llen == 0 {
		langs = defaultLangs
		llen = len(langs)
	}
	for i, l := range langs {
		lang += l
		if i < llen-1 {
			lang += "+"
		}
	}
	threads := c.TesseractThreads
	if threads > maxThreads {
		threads = maxThreads
	} else if threads < minThreads {
		threads = minThreads
	}
	args := []string{"-l", lang, "--oem", fmt.Sprintf("%d", c.TesseractOEM)}
	s = &Scanner{
		log,
		c.DataDir,
		time.Second * time.Duration(c.CMDTimeout),
		Tesseract{tcmd, args, []string{fmt.Sprintf("OMP_THREAD_LIMIT=%d", threads)}},
		PDF{pcmd, []string{"-layout", "-nopgbrk", "-eol", "unix"}},
	}
	return
}

func (s *Scanner) Scan(file string) (filename, text string, err error) {
	log := s.log
	file = fs.Clean(file)
	scanfile := filepath.Join(s.base, file)
	var pdffile string
	ft := s.newFileType(scanfile)
	if !ft.pdf {
		if !ft.image {
			if ft.text {
				log.V(2).Info("indexing plaintext file", "file", scanfile)
				var buf []byte
				buf, err = os.ReadFile(scanfile)
				if err != nil {
					return
				}
				filename = file
				text = string(nlr.ReplaceAll(wsr.ReplaceAll(buf, ws), nl))
				return
			} else {
				err = errors.New("error: input: no image, pdf or text file")
				return
			}
		}
		var (
			pdf []byte
			f   *os.File
		)
		p, _ := filepath.Split(file)
		pdfCmd := exec.Command(s.t.cmd, append(s.t.args, []string{scanfile, "stdout", "pdf"}...)...)
		pdfCmd.Env = s.t.env
		log.V(2).Info("tesseract command", "env", pdfCmd.Env, "path", pdfCmd.Path, "args", pdfCmd.Args)
		pdf, err = timeout(pdfCmd, s.timeout)
		if err != nil {
			return
		}
		filename = filepath.Join(p, core.FileName()+pdfExt)
		pdffile = filepath.Join(s.base, filename)
		f, err = os.OpenFile(pdffile, os.O_RDWR|os.O_CREATE, fs.FileMode)
		if err != nil {
			f.Close()
			return
		}
		_, err = f.Write(pdf)
		if err != nil {
			f.Close()
			return
		}
		f.Close()
	} else {
		pdffile = scanfile
		filename = file
	}
	txtCmd := exec.Command(s.p.cmd, append(s.p.args, []string{pdffile, "-"}...)...)
	log.V(2).Info("text2pdf command", "path", txtCmd.Path, "args", txtCmd.Args)
	txt, err := timeout(txtCmd, s.timeout)
	if err != nil {
		return
	}
	if len(txt) <= 2 {
		err = errors.New("pdf does not contain any text information")
		return
	}
	text = string(nlr.ReplaceAll(wsr.ReplaceAll(txt, ws), nl))
	if !ft.pdf {
		err = os.Remove(scanfile)
	}
	return
}