go-coreutils/hashsum.go

package coreutils

import (
	"bufio"
	"crypto/md5"
	"errors"
	"fmt"
	"hash"
	"io"
	"os"
	"regexp"

	"github.com/alexflint/go-arg"
)

var MD5Regex = regexp.MustCompile("^(?P<hash>[0-9a-f]{32})  (?P<filename>.*)$")

func MD5Sum(r io.Reader) (hash.Hash, error) {
	h := md5.New()
	if _, err := io.Copy(h, r); err != nil {
		return nil, err
	}
	return h, nil
}

type CheckingResults struct {
	ImproperlyFormattedCount uint
	InvalidChecksumCount     uint
	FilesNotRead             uint
}

// SumFunc is a type of function that computes a hash.Hash for data in io.Reader
type SumFunc func(io.Reader) (hash.Hash, error)

// ImproperlyFormattedErr is an error return when a line for checking
// has an incorrect number or set of characters
var ImproperlyFormattedErr = errors.New("improperly formatted line")

func printHash(h hash.Hash, filename string) {
	fmt.Printf("%x  %v\n", h.Sum(nil), filename)
}

// PrintSumForFile prints checksum computed by f for file with filename
func PrintSumForFile(filename string, f SumFunc) error {
	file, err := os.Open(filename)
	if err != nil {
		return err
	}
	defer file.Close()

	h, err := f(file)
	if err != nil {
		return err
	}

	printHash(h, filename)
	return nil
}

// PrintSumForStdin prints checksum computed by f for data from stdin
func PrintSumForStdin(f SumFunc) {
	h, err := f(os.Stdin)
	if err != nil {
		PrintToStderr(err)
		os.Exit(1)
	}
	printHash(h, "<stdin>")
	os.Exit(0)
}

func checkLine(line string, re *regexp.Regexp, f SumFunc) (valid bool, err error) {
	matches := re.MatchString(line)
	if !matches {
		err = ImproperlyFormattedErr
		valid = false
		return
	}
	submatches := re.FindStringSubmatch(line)
	hash := submatches[1]
	filename := submatches[2]

	file, err := os.Open(filename)
	if err != nil {
		valid = false
		return
	}
	defer file.Close()
	h, err := f(file)
	if err != nil {
		valid = false
		return
	}

	result := "OK"
	valid = true
	if fmt.Sprintf("%x", h.Sum(nil)) != hash {
		result = "FAILED"
		valid = false
	}
	fmt.Printf("%s: %s\n", filename, result)

	return
}

// CheckSumsInReader checks checksums in r using f
func CheckSumsInReader(r io.Reader, re *regexp.Regexp, f SumFunc) (*CheckingResults, error) {
	scanner := bufio.NewScanner(r)

	results := CheckingResults{}
	for scanner.Scan() {
		valid, err := checkLine(scanner.Text(), re, f)
		if err == ImproperlyFormattedErr {
			results.ImproperlyFormattedCount++
		} else if err != nil {
			PrintToStderr(err)
			results.FilesNotRead++
		} else if !valid {
			results.InvalidChecksumCount++
		}
	}
	return &results, scanner.Err()
}

// CheckSumsInFile checks checksums in file with filename using f
func CheckSumsInFile(filename string, re *regexp.Regexp, f SumFunc) (*CheckingResults, error) {
	file, err := os.Open(filename)
	if err != nil {
		return nil, err
	}
	defer file.Close()

	return CheckSumsInReader(file, re, f)
}

// PrintCheckingResults prints the number of mismatched checksums and improperly formatted lines to os.Stderr
func PrintCheckingResults(results CheckingResults) {
	if results.InvalidChecksumCount > 0 {
		PrintToStderr(
			fmt.Sprintf("%v computed checksums did NOT match", results.InvalidChecksumCount),
		)
	}
	if results.ImproperlyFormattedCount > 0 {
		PrintToStderr(
			fmt.Sprintf("%v lines are improperly formatted", results.ImproperlyFormattedCount),
		)
	}
	if results.FilesNotRead > 0 {
		PrintToStderr(
			fmt.Sprintf("%v listed files could not be read", results.FilesNotRead),
		)
	}
}

// SumMain is a function you're supposed to run in main of a hash sum utility
func SumMain(re *regexp.Regexp, f SumFunc) {
	var args struct {
		Check bool     `arg:"-c"`
		Files []string `arg:"positional"`
	}

	arg.MustParse(&args)

	exitCode := 0
	if len(args.Files) == 0 {
		if !args.Check {
			PrintSumForStdin(f)
		} else {
			results, err := CheckSumsInReader(os.Stdin, re, f)
			if err != nil {
				PrintToStderr(err)
				exitCode = 1
			}
			PrintCheckingResults(*results)
		}
	}

	totalResults := CheckingResults{}
	for _, filename := range args.Files {
		if !args.Check {
			err := PrintSumForFile(filename, f)
			if err != nil {
				PrintToStderr(err)
				exitCode = 1
			}
		} else {
			results, err := CheckSumsInFile(filename, re, f)
			if err != nil {
				PrintToStderr(err)
				exitCode = 1
			}
			if results != nil {
				totalResults.ImproperlyFormattedCount += results.ImproperlyFormattedCount
				totalResults.InvalidChecksumCount += results.InvalidChecksumCount
			}
		}
	}
	PrintCheckingResults(totalResults)
	os.Exit(exitCode)
}