~robertgzr/corrupter

0e0d962637a8b42e555d2ae2364722fbc24b3e9d — Robert Günzler 1 year, 1 month ago 9b24839
Add screenlocker

Signed-off-by: Robert Günzler <r@gnzler.io>
3 files changed, 321 insertions(+), 0 deletions(-)

A cmd/lock/lock.go
M go.mod
A go.sum
A cmd/lock/lock.go => cmd/lock/lock.go +317 -0
@@ 0,0 1,317 @@
// vim: ft=go
package main

import (
	"bufio"
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"image"
	"image/png"
	"io"
	"log"
	"math/rand"
	"os"
	"os/exec"
	"sync"
	"time"

	"golang.org/x/sys/unix"
)

const (
	mainColor  = "#4E4371"
	clearColor = "#AFD700"
	verColor   = "#5F8787"
	wrongColor = "#AF005F"
)

const (
	mag       = 7.
	bheight   = 10
	boffset   = 30.
	strideC   = .1
	lag       = .005
	lrC       = -7.
	lgC       = 0.
	lbC       = 3.
	stdOffset = 10
	add       = 10.
	meanAbber = 10.
	stdAbber  = 10.
)

func main() {
	log.SetOutput(os.Stderr)
	log.SetFlags(log.Ltime)

	swaylock := exec.Command("swaylock",
		"--color=#000000EE",
		"--daemonize",
		"--font-size=16",
		"--font=Iosevka Aile",
		"--hide-keyboard-layout",
		"--ignore-empty-password",
		"--indicator-caps-lock",
		"--indicator-idle-visible",
		"--inside-clear-color=#000000BB",
		"--inside-color=#000000BB",
		"--inside-ver-color=#000000BB",
		"--inside-wrong-color="+wrongColor+"BB",
		"--key-hl-color=#FFFFFF55",
		"--ring-clear-color="+clearColor+"EE",
		"--ring-color="+mainColor+"EE",
		"--ring-ver-color="+verColor+"EE",
		"--ring-wrong-color="+wrongColor+"EE",
		"--scaling=fill",
		"--show-failed-attempts",
		"--text-clear-color=#FFFFFFFF",
		"--text-color=#FFFFFFFF",
		"--text-ver-color=#FFFFFFFF")

	var wg sync.WaitGroup

	{
		swqOutput, err := exec.Command("swaymsg", "--type", "get_tree", "--raw").Output()
		if err != nil {
			log.Fatal(err)
		}
		var swqJson struct {
			Nodes []struct {
				Name string
			}
		}
		if err := json.Unmarshal(swqOutput, &swqJson); err != nil {
			log.Fatal(err)
		}

		for _, output := range swqJson.Nodes {
			if output.Name == "__i3" {
				continue
			}
			wg.Add(1)
			go func(output string) {
				defer wg.Done()

				grim := exec.Command("grim", "-o", output, "-")
				rc, err := grim.StdoutPipe()
				if err != nil {
					log.Fatalf("%s: %v", output, err)
				}
				defer rc.Close()

				// capture screenshot
				log.Printf("+ %v", grim.Args)
				if err := grim.Start(); err != nil {
					log.Fatalf("%s: %v", output, err)
				}

				var buf bytes.Buffer
				defer buf.Reset()

				// process screenshot
				log.Printf("%s: processing", output)
				if err := corrupter(rc, &buf, time.Now().Unix()); err != nil {
					log.Fatalf("%s: %v", output, err)
				}

				log.Printf("%s: processed %d bytes", output, buf.Len())

				// memfd to pass image from internal processing
				// to swaylock
				fd, err := unix.MemfdCreate("lock-"+output, 0)
				if err != nil {
					log.Fatalf("%s: %v", output, err)
				}
				if err := unix.Ftruncate(fd, int64(buf.Len())); err != nil {
					log.Fatal(err)
				}
				mem, err := unix.Mmap(fd, 0, buf.Len(), unix.PROT_READ|unix.PROT_WRITE, unix.MAP_SHARED)
				if err != nil {
					log.Fatalf("%s: %v", output, err)
				}
				copy(mem, buf.Bytes())
				if err := unix.Munmap(mem); err != nil {
					log.Fatalf("%s: %v", output, err)
				}

				swaylock.ExtraFiles = append(swaylock.ExtraFiles, os.NewFile(uintptr(fd), fmt.Sprintf("/proc/self/fd/%d", fd)))
				swaylock.Args = append(swaylock.Args, "--image="+output+":"+fmt.Sprintf("/proc/self/fd/%d", 2+len(swaylock.ExtraFiles)))
			}(output.Name)
		}
	}

	{
		wg.Add(1)
		go func() {
			defer wg.Done()
			runparts := exec.Command("run-parts", "-v", os.ExpandEnv("${XDG_DATA_HOME}/lock.d"))
			logCmd(runparts)
			if err := runparts.Run(); err != nil {
				log.Fatal(err)
			}
		}()
	}

	wg.Wait()

	logCmd(swaylock)
	if err := swaylock.Run(); err != nil {
		log.Fatal(err)
	}
}

func logCmd(c *exec.Cmd) {
	handler := func(pipe io.ReadCloser, err error) {
		if err != nil {
			log.Fatalf("%s: %v", c.Args[0], err)
		}
		defer pipe.Close()
		scan := bufio.NewScanner(pipe)
		scan.Split(bufio.ScanLines)
		for scan.Scan() {
			log.Printf(scan.Text())
		}
	}
	log.Printf("+ %+v", c.Args)
	go handler(c.StdoutPipe())
	go handler(c.StderrPipe())
}

func corrupter(reader io.Reader, writer io.Writer, seed int64) error {
	seededRand := rand.New(rand.NewSource(seed))

	// flag.Args() contain all non-option arguments, i.e., our input and output files
	m, err := png.Decode(reader)
	if err != nil {
		return err
	}

	// trying to obtain raw pointers to color data, since .At(), .Set() are very slow
	m_raw_stride, m_raw_pix := 0, []uint8(nil)

	switch m.(type) {
	default:
		return errors.New("unknown image type")
	case *image.NRGBA:
		m_raw := m.(*image.NRGBA)
		m_raw_stride = m_raw.Stride
		m_raw_pix = m_raw.Pix
	case *image.RGBA:
		m_raw := m.(*image.RGBA)
		m_raw_stride = m_raw.Stride
		m_raw_pix = m_raw.Pix
	}

	b := m.Bounds()

	// first stage is dissolve+block corruption
	new_img := image.NewNRGBA(b)
	line_off := 0
	stride := 0.
	yset := 0
	for y := 0; y < b.Max.Y; y++ {
		for x := 0; x < b.Max.X; x++ {
			// Every BHEIGHT lines in average a new distorted block begins
			if seededRand.Intn(bheight*b.Max.X) == 0 {
				line_off = offset(seededRand, boffset)
				stride = seededRand.NormFloat64() * strideC
				yset = y
			}
			// at the line where the block has begun, we don't want to offset the image
			// so stride_off is 0 on the block's line
			stride_off := int(stride * float64(y-yset))

			// offset is composed of the blur, block offset, and skew offset (stride)
			offx := offset(seededRand, mag) + line_off + stride_off
			offy := offset(seededRand, mag)

			// copy the corresponding pixel (4 bytes) to the new image
			src_idx := m_raw_stride*wrap(y+offy, b.Max.Y) + 4*wrap(x+offx, b.Max.X)
			dst_idx := new_img.Stride*y + 4*x

			copy(new_img.Pix[dst_idx:dst_idx+4], m_raw_pix[src_idx:src_idx+4])
		}
	}

	// second stage is adding per-channel scan inconsistency and brightening
	new_img1 := image.NewNRGBA(b)

	lr, lg, lb := lrC, lgC, lbC
	for y := 0; y < b.Max.Y; y++ {
		for x := 0; x < b.Max.X; x++ {
			lr += seededRand.NormFloat64() * lag
			lg += seededRand.NormFloat64() * lag
			lb += seededRand.NormFloat64() * lag
			offx := offset(seededRand, stdOffset)

			// obtain source pixel base offsets. red/blue border is also smoothed by offx
			ra_idx := new_img.Stride*y + 4*wrap(x+int(lr)-offx, b.Max.X)
			g_idx := new_img.Stride*y + 4*wrap(x+int(lg), b.Max.X)
			b_idx := new_img.Stride*y + 4*wrap(x+int(lb)+offx, b.Max.X)

			// pixels are stored in (r, g, b, a) order in memory
			r := new_img.Pix[ra_idx]
			a := new_img.Pix[ra_idx+3]
			g := new_img.Pix[g_idx+1]
			b := new_img.Pix[b_idx+2]

			r, g, b = brighten(r, uint8(add)), brighten(g, uint8(add)), brighten(b, uint8(add))

			// copy the corresponding pixel (4 bytes) to the new image
			dst_idx := new_img1.Stride*y + 4*x
			copy(new_img1.Pix[dst_idx:dst_idx+4], []uint8{r, g, b, a})
		}
	}

	// third stage is to add chromatic abberation+chromatic trails
	// (trails happen because we're changing the same image we process)
	for y := 0; y < b.Max.Y; y++ {
		for x := 0; x < b.Max.X; x++ {
			offx := meanAbber + offset(seededRand, stdAbber) // lower offset arg = longer trails

			// obtain source pixel base offsets. only red and blue are distorted
			ra_idx := new_img1.Stride*y + 4*wrap(x+offx, b.Max.X)
			g_idx := new_img1.Stride*y + 4*x
			b_idx := new_img1.Stride*y + 4*wrap(x-offx, b.Max.X)

			// pixels are stored in (r, g, b, a) order in memory
			r := new_img1.Pix[ra_idx]
			a := new_img1.Pix[ra_idx+3]
			g := new_img1.Pix[g_idx+1]
			b := new_img1.Pix[b_idx+2]

			// copy the corresponding pixel (4 bytes) to the SAME image. this gets us nice colorful trails
			dst_idx := new_img1.Stride*y + 4*x
			copy(new_img1.Pix[dst_idx:dst_idx+4], []uint8{r, g, b, a})
		}
	}

	// write the image
	e := png.Encoder{CompressionLevel: png.NoCompression}
	return e.Encode(writer, new_img1)
}

// force x to stay in [0, b) range. x is assumed to be in [-b,2*b) range
func wrap(x, b int) int {
	if x < 0 {
		return x + b
	}
	if x >= b {
		return x - b
	}
	return x
}

// get normally distributed (rounded to int) value with the specified std. dev.
func offset(seededRand *rand.Rand, stddev float64) int {
	sample := seededRand.NormFloat64() * stddev
	return int(sample)
}

// brighten the color safely, i.e., by simultaneously reducing contrast
func brighten(r uint8, add uint8) uint8 {
	r32, add32 := uint32(r), uint32(add)
	return uint8(r32 - r32*add32/255 + add32)
}

M go.mod => go.mod +2 -0
@@ 1,3 1,5 @@
module github.com/r00tman/corrupter

go 1.12

require golang.org/x/sys v0.4.0 // indirect

A go.sum => go.sum +2 -0
@@ 0,0 1,2 @@
golang.org/x/sys v0.4.0 h1:Zr2JFtRQNX3BCZ8YtxRE9hNJYC8J6I1MVbMg6owUp18=
golang.org/x/sys v0.4.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=