1package hints
2
3import (
4 "bufio"
5 "encoding/json"
6 "fmt"
7 "image"
8 "image/color"
9 "image/png"
10 "os"
11 "path/filepath"
12 "sort"
13 "strings"
14 "syscall"
15 "time"
16
17 "github.com/bnema/neferwl/internal/ports"
18)
19
20const (
21 // maxDumps is the number of dump groups kept in the directory.
22 maxDumps = 50
23 outlineSize = 2
24)
25
26// Options tunes Run.
27type Options struct {
28 // DumpDir, when not empty, receives the scanned frames (PNG, annotated
29 // PNG and JSON). They contain the screen in clear.
30 DumpDir string
31 // Source, when not nil, finds targets in the accessibility tree of the
32 // windows on screen; OCR reads the rest.
33 Source targetSource
34}
35
36// dumpJob is one scan to write. pix is the capture buffer of the scan, owned
37// by the job: writeDump converts it in place.
38type dumpJob struct {
39 when time.Time
40 req ports.HintScanRequest
41 pix []byte
42 m *mapping // owns pix; released after the dump
43 stride int
44 targets []ports.HintTarget
45 crops []image.Rectangle
46 words []OCRWord
47 // atspi is the number of leading targets that come from the
48 // accessibility tree; the others come from OCR.
49 atspi int
50 captureD time.Duration
51 scanD time.Duration
52 atspiD time.Duration
53}
54
55type dumpTarget struct {
56 Kind string `json:"kind"`
57 Source string `json:"source"`
58 Rect [4]int `json:"rect"`
59}
60
61type dumpMeta struct {
62 Output string `json:"output"`
63 Width int `json:"width"`
64 Height int `json:"height"`
65 Stride int `json:"stride"`
66 Scale float64 `json:"scale"`
67 LinksOnly bool `json:"links_only"`
68 CaptureMS float64 `json:"capture_ms"`
69 ScanMS float64 `json:"scan_ms"`
70 AtspiMS float64 `json:"atspi_ms"`
71 Crops []image.Rectangle `json:"ocr_crops"`
72 Words []OCRWord `json:"words"`
73 Targets []dumpTarget `json:"targets"`
74}
75
76func kindName(k ports.HintKind) string {
77 switch k {
78 case ports.HintLink:
79 return "link"
80 case ports.HintButton:
81 return "button"
82 default:
83 return "text"
84 }
85}
86
87func kindColor(k ports.HintKind) color.RGBA {
88 switch k {
89 case ports.HintLink:
90 return color.RGBA{0x00, 0xc8, 0x53, 0xff}
91 case ports.HintButton:
92 return color.RGBA{0x29, 0x62, 0xff, 0xff}
93 default:
94 return color.RGBA{0x80, 0x80, 0x80, 0xff}
95 }
96}
97
98// atspiColor is the outline of targets found in the accessibility tree.
99var atspiColor = color.RGBA{0xff, 0x00, 0xff, 0xff}
100
101func ms(d time.Duration) float64 { return float64(d.Microseconds()) / 1000 }
102
103// writeDump writes the three files of j into dir, then prunes old groups.
104func writeDump(dir string, j *dumpJob) error {
105 if err := secureDir(dir); err != nil {
106 return err
107 }
108 prefix := filepath.Join(dir, j.when.Format("20060102-150405.000")+"-"+safeName(j.req.Output))
109
110 // The frame is converted in place to RGBA and encoded straight from the
111 // capture buffer: no copy. The job owns the buffer, which is released
112 // with it.
113 for y := 0; y < j.req.Height; y++ {
114 row := j.pix[y*j.stride : y*j.stride+j.req.Width*4]
115 for x := 0; x < len(row); x += 4 {
116 row[x], row[x+2], row[x+3] = row[x+2], row[x], 0xff
117 }
118 }
119 img := &image.RGBA{Pix: j.pix, Stride: j.stride, Rect: image.Rect(0, 0, j.req.Width, j.req.Height)}
120 if err := writePNG(prefix+".png", img); err != nil {
121 return err
122 }
123
124 scale := j.req.Scale
125 if scale <= 0 {
126 scale = 1
127 }
128 for _, r := range j.crops {
129 outline(img, r, color.RGBA{0x80, 0x80, 0x80, 0xff})
130 }
131 for i, t := range j.targets {
132 r := image.Rect(int(float64(t.Rect.X)*scale), int(float64(t.Rect.Y)*scale),
133 int(float64(t.Rect.X+t.Rect.W)*scale), int(float64(t.Rect.Y+t.Rect.H)*scale))
134 c := kindColor(t.Kind)
135 if i < j.atspi {
136 c = atspiColor
137 }
138 outline(img, r, c)
139 }
140 if err := writePNG(prefix+".annotated.png", img); err != nil {
141 return err
142 }
143
144 meta := dumpMeta{
145 Output: j.req.Output, Width: j.req.Width, Height: j.req.Height, Stride: j.stride,
146 Scale: j.req.Scale, LinksOnly: j.req.LinksOnly,
147 CaptureMS: ms(j.captureD), ScanMS: ms(j.scanD), AtspiMS: ms(j.atspiD),
148 Crops: j.crops, Words: j.words,
149 Targets: make([]dumpTarget, 0, len(j.targets)),
150 }
151 for i, t := range j.targets {
152 src := "ocr"
153 if i < j.atspi {
154 src = "atspi"
155 }
156 meta.Targets = append(meta.Targets, dumpTarget{Kind: kindName(t.Kind), Source: src, Rect: [4]int{t.Rect.X, t.Rect.Y, t.Rect.X + t.Rect.W, t.Rect.Y + t.Rect.H}})
157 }
158 data, err := json.Marshal(meta)
159 if err != nil {
160 return err
161 }
162 if err := writeFile(prefix+".json", func(f *os.File) error { _, err := f.Write(append(data, '\n')); return err }); err != nil {
163 return err
164 }
165 return prune(dir, maxDumps)
166}
167
168func safeName(s string) string {
169 return strings.Map(func(r rune) rune {
170 if r == '/' || r == 0 {
171 return '_'
172 }
173 return r
174 }, s)
175}
176
177func writePNG(path string, img image.Image) error {
178 enc := png.Encoder{CompressionLevel: png.BestSpeed}
179 return writeFile(path, func(f *os.File) error {
180 bw := bufio.NewWriterSize(f, 1<<16)
181 if err := enc.Encode(bw, img); err != nil {
182 return err
183 }
184 return bw.Flush()
185 })
186}
187
188// secureDir makes sure dir is a real directory owned by the current user and
189// private to it, whatever it was before.
190func secureDir(dir string) error {
191 if err := os.MkdirAll(dir, 0o700); err != nil {
192 return err
193 }
194 fi, err := os.Lstat(dir)
195 if err != nil {
196 return err
197 }
198 if !fi.IsDir() {
199 return fmt.Errorf("dump dir %s is not a directory", dir)
200 }
201 if st, ok := fi.Sys().(*syscall.Stat_t); ok && int(st.Uid) != os.Getuid() {
202 return fmt.Errorf("dump dir %s is not owned by the current user", dir)
203 }
204 return os.Chmod(dir, 0o700)
205}
206
207// writeFile fills a fresh private temporary file in the directory of path and
208// renames it over path: an existing file or symlink there is replaced, never
209// followed.
210func writeFile(path string, fill func(*os.File) error) error {
211 f, err := os.CreateTemp(filepath.Dir(path), ".dump-*")
212 if err != nil {
213 return err
214 }
215 tmp := f.Name()
216 err = f.Chmod(0o600)
217 if err == nil {
218 err = fill(f)
219 }
220 if cerr := f.Close(); err == nil {
221 err = cerr
222 }
223 if err == nil {
224 err = os.Rename(tmp, path)
225 }
226 if err != nil {
227 _ = os.Remove(tmp)
228 }
229 return err
230}
231
232// outline draws a 2 px frame inside r, clipped to img.
233func outline(img *image.RGBA, r image.Rectangle, c color.RGBA) {
234 r = r.Intersect(img.Rect)
235 if r.Empty() {
236 return
237 }
238 for y := r.Min.Y; y < r.Max.Y; y++ {
239 for x := r.Min.X; x < r.Max.X; x++ {
240 if x-r.Min.X < outlineSize || r.Max.X-1-x < outlineSize || y-r.Min.Y < outlineSize || r.Max.Y-1-y < outlineSize {
241 img.SetRGBA(x, y, c)
242 }
243 }
244 }
245}
246
247// groupOf returns the "<timestamp>-<output>" prefix of a dump file name.
248func groupOf(name string) string {
249 for _, suf := range []string{".annotated.png", ".png", ".json"} {
250 if strings.HasSuffix(name, suf) {
251 return strings.TrimSuffix(name, suf)
252 }
253 }
254 return ""
255}
256
257// prune keeps the keep most recent groups of dir; timestamps sort by name.
258func prune(dir string, keep int) error {
259 entries, err := os.ReadDir(dir)
260 if err != nil {
261 return err
262 }
263 groups := map[string][]string{}
264 for _, e := range entries {
265 if g := groupOf(e.Name()); g != "" && !e.IsDir() {
266 groups[g] = append(groups[g], e.Name())
267 }
268 }
269 if len(groups) <= keep {
270 return nil
271 }
272 names := make([]string, 0, len(groups))
273 for g := range groups {
274 names = append(names, g)
275 }
276 sort.Strings(names)
277 for _, g := range names[:len(names)-keep] {
278 for _, f := range groups[g] {
279 if err := os.Remove(filepath.Join(dir, f)); err != nil {
280 return fmt.Errorf("prune: %w", err)
281 }
282 }
283 }
284 return nil
285}