tools/syz-usbgen/usbgen.go

// Copyright 2019 syzkaller project authors. All rights reserved.
// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.

package main

import (
	"encoding/hex"
	"flag"
	"fmt"
	"os"
	"regexp"
	"sort"

	"github.com/google/syzkaller/pkg/osutil"
	"github.com/google/syzkaller/pkg/tool"
)

func main() {
	flag.Parse()
	args := flag.Args()
	if len(args) != 2 {
		usage()
	}

	syslog, err := os.ReadFile(args[0])
	if err != nil {
		tool.Failf("failed to read file %v: %v", args[0], err)
	}

	usbIds := extractIds(syslog, "USBID", 34)
	hidIds := extractIds(syslog, "HIDID", 24)

	output := []byte(`// Code generated by tools/syz-usbgen. DO NOT EDIT.
// See docs/linux/external_fuzzing_usb.md

package linux

`)
	output = append(output, generateIdsVar(usbIds, "usbIds")...)
	output = append(output, []byte("\n")...)
	output = append(output, generateIdsVar(hidIds, "hidIds")...)

	if err := osutil.WriteFile(args[1], output); err != nil {
		tool.Failf("failed to output file %v: %v", args[1], err)
	}
}

func extractIds(syslog []byte, prefix string, size int) map[string][]string {
	re := fmt.Sprintf("%s: ([0-9a-f]{%d}) \\((.+)\\)", prefix, size)
	r := regexp.MustCompile(re)
	matches := r.FindAllSubmatch(syslog, -1)
	// Map from matching substring to capture groups.
	uniqueMatches := make(map[string][][]byte)
	for _, match := range matches {
		uniqueMatches[string(match[0])] = match[1:]
	}
	// Map from driver name to slice of USB IDs.
	driverIDs := make(map[string][]string, 0)
	for _, groups := range uniqueMatches {
		id := string(groups[0])
		driver := string(groups[1])
		driverIDs[driver] = append(driverIDs[driver], id)
	}
	// Keep IDs sorted for consistent output between runs.
	for driver := range driverIDs {
		sort.Strings(driverIDs[driver])
	}
	return driverIDs
}

func generateIdsVar(driverIDs map[string][]string, name string) []byte {
	// Sort driver names for consistent output between runs.
	drivers := make([]string, 0, len(driverIDs))
	for driver := range driverIDs {
		drivers = append(drivers, driver)
	}
	sort.Strings(drivers)

	// Generate a map variable that stores USB IDs for each driver.
	totalIDs := 0
	output := []byte(fmt.Sprintf("var %s = map[string]string{\n", name))
	for _, driver := range drivers {
		ids := driverIDs[driver]
		outputDriver := fmt.Sprintf("\t\"%s\": ", driver)
		output = append(output, []byte(outputDriver)...)
		for i, id := range ids {
			decodedID, err := hex.DecodeString(id)
			if err != nil {
				tool.Failf("failed to decode hex string %v: %v", id, err)
			}
			prefix := "\t\t"
			suffix := " +"
			if i == 0 {
				prefix = ""
			}
			if i == len(ids)-1 {
				suffix = ","
			}
			outputID := fmt.Sprintf("%v%#v%v\n", prefix, string(decodedID), suffix)
			output = append(output, []byte(outputID)...)
		}
		totalIDs += len(ids)
	}
	output = append(output, []byte("}\n\n")...)

	// Generate a variable that stores all USB IDs together.
	output = append(output, []byte(fmt.Sprintf("var %sAll = ", name))...)
	for i, driver := range drivers {
		prefix := "\t"
		suffix := " +"
		if i == 0 {
			prefix = ""
		}
		if i == len(drivers)-1 {
			suffix = ""
		}
		outputDriver := fmt.Sprintf("%v%s[\"%s\"]%v\n", prefix, name, driver, suffix)
		output = append(output, []byte(outputDriver)...)
	}
	if len(drivers) == 0 {
		output = append(output, []byte("\"\"")...)
	}

	fmt.Printf("%v %s ids written\n", totalIDs, name)

	return output
}

func usage() {
	fmt.Fprintf(os.Stderr, "usage:\n")
	fmt.Fprintf(os.Stderr, "  syz-usbgen syslog.txt sys/linux/init_vusb_ids.go\n")
	os.Exit(1)
}