usblocker/usblocker.go

package main

import (
	"fmt"
	"log"
	"os"
	_ "os/exec"
	"path/filepath"
	"strings"
	"time"
)

/* Shows the help message
 */
func ShowHelp() {
	fmt.Println("Watches a device and locks the screen if the device is missing.")
	fmt.Println()
	fmt.Println("Usage:")
	fmt.Println("  usblocker run:     Runs USBLocker in foreground.")
	//fmt.Println("  usblocker daemon:  Runs in background.")
	//fmt.Println("  usblocker create:  Prepares a device for usage with USBLocker.")
	//fmt.Println("                     (Not yet implemented!)")
}

/* The watcher himself!
 */
func WatcherUnlocked(device string) {
	// Watch!
	for {
		_, errs := os.Stat(device)
		if errs != nil {
			break
		}
		time.Sleep(time.Duration(1) * time.Second)
	}
	log.Println("Locked")
	var args []string
	args = append(args, "-nolock")
	//args = append(args, "-info")
	//args = append(args, "")
	args = append(args, "-mode")
	args = append(args, "blank")
	//log.Println(args)
	var attr os.ProcAttr
	lpid, err := os.StartProcess("/usr/bin/xlock", args, &attr)
	if err != nil {
		log.Fatal("Error: ", err)
	}
	WatcherLocked(lpid)
}

/* Watcher for locked state / no device
 */
func WatcherLocked(lpid *os.Process) {
	var passphrase []byte // TEH KEY!!1!
	var devices []string  // list of devices
	var lfile *os.File    // device file
	ldev := ""            // device file name
	version := "v0"       // version string

	// Set the passphrase
	for i := range passphrase {
		passphrase[i] = 0
	}

	for {
		var err error
		devices, err = filepath.Glob("/dev/?d?1")
		if err != nil {
			log.Fatal("Error while getting list of aviable devices (/dev/?d?1):", err)
		}

		// Search lock device
		for i := range devices {
			log.Println("Checking", devices[i], "for usblocker data...")
			lfile, err = os.Open(devices[i])
			if err != nil {
				log.Fatal("Error while opening the device file:", err)
			}
			// Check for magic string
			magictest := []byte("         ")
			length, err := lfile.Read(magictest)
			if err != nil {
				log.Fatal("Error while reading the device file:", err)
				_ = lfile.Close()
			}
			if length == 9 && string(magictest) == "usblocker" {
				ldev = devices[i]
				break
			}
		}

		// Check lockdev version
		if ldev != "" {
			log.Println("Checking", ldev, "for right version...")
			versiontest := []byte("       ")
			length, err := lfile.Read(versiontest)
			_ = lfile.Close()
			if err != nil {
				log.Fatal("Error while reading the device file:", err)
			}
			if length == 7 && strings.HasPrefix(string(versiontest), version) {
				// And finally: UNLOCK!
				log.Println("Unlocked")
				if lpid != nil {
					lpid.Kill()
					_, _ = lpid.Wait()
				}
				WatcherUnlocked(ldev)
			}
		}

		time.Sleep(time.Duration(1) * time.Second)
	}
}

/* The main function
 */
func main() {
	// Determining what to do
	if len(os.Args) == 1 {
		ShowHelp()
		return
	}
	switch os.Args[1] {
	case "run":
	case "daemon":
		log.Fatal("Not yet implemented.")
	case "create":
		log.Fatal("Not yet implemented.")
	default:
		ShowHelp()
	}

	// Check for root
	if os.Getuid() != 0 {
		log.Fatal("YU NO GOT ROOT??")
	}

	// Greet the user :)
	//log.Println("USBLocker", version)

	WatcherLocked(nil)
}