main.go

   1 /*
   2 glocate -- ZFS-diff-friendly locate-like utility
   3 Copyright (C) 2022 Sergey Matveev <stargrave@stargrave.org>
   4
   5 This program is free software: you can redistribute it and/or modify
   6 it under the terms of the GNU General Public License as published by
   7 the Free Software Foundation, version 3 of the License.
   8
   9 This program is distributed in the hope that it will be useful,
  10 but WITHOUT ANY WARRANTY; without even the implied warranty of
  11 MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  12 GNU General Public License for more details.
  13
  14 You should have received a copy of the GNU General Public License
  15 along with this program.  If not, see <http://www.gnu.org/licenses/>.
  16 */
  17
  18 package main
  19
  20 import (
  21         "bufio"
  22         "encoding/gob"
  23         "flag"
  24         "fmt"
  25         "io"
  26         "io/fs"
  27         "log"
  28         "os"
  29         "path"
  30         "sort"
  31         "strconv"
  32         "strings"
  33         "syscall"
  34         "time"
  35
  36         "github.com/dustin/go-humanize"
  37         "github.com/klauspost/compress/zstd"
  38 )
  39
  40 type File struct {
  41         Name  string
  42         Size  uint64
  43         Mtime int64
  44         Files []File
  45 }
  46
  47 type ByName []File
  48
  49 func (a ByName) Len() int {
  50         return len(a)
  51 }
  52
  53 func (a ByName) Swap(i, j int) {
  54         a[i], a[j] = a[j], a[i]
  55 }
  56
  57 func (a ByName) Less(i, j int) bool {
  58         return a[i].Name < a[j].Name
  59 }
  60
  61 func (file *File) IsDir() bool {
  62         return file.Name[len(file.Name)-1] == '/'
  63 }
  64
  65 func walk(root string) ([]File, uint64, error) {
  66         fd, err := os.Open(root)
  67         if err != nil {
  68                 return nil, 0, err
  69         }
  70         var files []File
  71         var size uint64
  72         var info fs.FileInfo
  73         for {
  74                 ents, err := fd.ReadDir(1 << 10)
  75                 if err != nil {
  76                         if err == io.EOF {
  77                                 break
  78                         }
  79                         fd.Close()
  80                         return nil, 0, err
  81                 }
  82                 for _, ent := range ents {
  83                         file := File{Name: ent.Name()}
  84                         fullPath := path.Join(root, file.Name)
  85                         if ent.IsDir() {
  86                                 file.Name += "/"
  87                         }
  88                         info, err = ent.Info()
  89                         if err != nil {
  90                                 log.Println("can not stat:", fullPath, ":", err)
  91                                 files = append(files, file)
  92                                 continue
  93                         }
  94                         file.Mtime = info.ModTime().Unix()
  95                         if ent.IsDir() {
  96                                 file.Files, file.Size, err = walk(fullPath)
  97                                 if err != nil {
  98                                         log.Println("can not walk:", fullPath, ":", err)
  99                                         files = append(files, file)
 100                                         continue
 101                                 }
 102                         } else if info.Mode().IsRegular() {
 103                                 file.Size = uint64(info.Size())
 104                         }
 105                         files = append(files, file)
 106                         size += file.Size
 107                 }
 108         }
 109         fd.Close()
 110         sort.Sort(ByName(files))
 111         return files, size, nil
 112 }
 113
 114 func usage() {
 115         log.Println("usage")
 116         os.Exit(1)
 117 }
 118
 119 func load(dbPath string) *File {
 120         fd, err := os.Open(dbPath)
 121         if err != nil {
 122                 log.Fatalln(err)
 123         }
 124         defer fd.Close()
 125         comp, err := zstd.NewReader(fd)
 126         if err != nil {
 127                 log.Fatalln(err)
 128         }
 129         dec := gob.NewDecoder(comp)
 130         var file File
 131         err = dec.Decode(&file)
 132         if err != nil {
 133                 log.Fatalln(err)
 134         }
 135         comp.Close()
 136         return &file
 137 }
 138
 139 func (db *File) dump(dbPath string) error {
 140         tmp, err := os.CreateTemp(path.Dir(dbPath), "glocate")
 141         if err != nil {
 142                 return err
 143         }
 144         defer os.Remove(tmp.Name())
 145         comp, err := zstd.NewWriter(
 146                 tmp, zstd.WithEncoderLevel(zstd.SpeedBestCompression),
 147         )
 148         if err != nil {
 149                 return err
 150         }
 151         enc := gob.NewEncoder(comp)
 152         err = enc.Encode(db)
 153         if err != nil {
 154                 return err
 155         }
 156         err = comp.Close()
 157         if err != nil {
 158                 return err
 159         }
 160         err = tmp.Close()
 161         if err != nil {
 162                 return err
 163         }
 164         umask := syscall.Umask(0)
 165         syscall.Umask(umask)
 166         err = os.Chmod(tmp.Name(), os.FileMode(0666&^umask))
 167         if err != nil {
 168                 return err
 169         }
 170         return os.Rename(tmp.Name(), dbPath)
 171 }
 172
 173 func (file *File) listBeauty(indent string, n int, isLast, veryFirst bool) {
 174         if veryFirst {
 175                 fmt.Printf("[%s]\n", humanize.IBytes(file.Size))
 176         } else {
 177                 var box string
 178                 if isLast {
 179                         box = "└"
 180                 } else {
 181                         box = "├"
 182                 }
 183                 name := file.Name
 184                 fmt.Printf("%s%s %s\t№%d [%s] %s\n",
 185                         indent, box, name, n, humanize.IBytes(file.Size),
 186                         time.Unix(file.Mtime, 0).Format("2006-01-02"),
 187                 )
 188                 if isLast {
 189                         indent += "  "
 190                 } else {
 191                         indent += "│ "
 192                 }
 193         }
 194         for n, f := range file.Files {
 195                 n++
 196                 f.listBeauty(indent, n, n == len(file.Files), false)
 197         }
 198 }
 199
 200 func (file *File) listSimple(root string, veryFirst bool) {
 201         name := file.Name
 202         fmt.Println(
 203                 strconv.FormatUint(file.Size, 10),
 204                 time.Unix(file.Mtime, 0).Format("2006-01-02T15:04:05"),
 205                 root+name,
 206         )
 207         if veryFirst {
 208                 name = ""
 209         }
 210         for _, f := range file.Files {
 211                 f.listSimple(root+name, false)
 212         }
 213 }
 214
 215 func (file *File) listFiles(root string, veryFirst bool) {
 216         name := file.Name
 217         if veryFirst {
 218                 root = ""
 219         } else {
 220                 fmt.Println(root + name)
 221                 root += name
 222         }
 223         for _, f := range file.Files {
 224                 f.listFiles(root, false)
 225         }
 226 }
 227
 228 func (db *File) find(p string) (file *File, parents []*File, idx int, err error) {
 229         file = db
 230         var f File
 231 Entities:
 232         for _, ent := range strings.Split(p, "/") {
 233                 for idx, f = range file.Files {
 234                         if (ent == f.Name) || (ent+"/" == f.Name) {
 235                                 parents = append(parents, file)
 236                                 file = &f
 237                                 continue Entities
 238                         }
 239                 }
 240                 err = fmt.Errorf("no entity found: %s", ent)
 241                 return
 242         }
 243         return
 244 }
 245
 246 func (db *File) remove(p string) error {
 247         file, parents, idx, err := db.find(p)
 248         if err != nil {
 249                 return err
 250         }
 251         lastParent := parents[len(parents)-1]
 252         lastParent.Files = append(
 253                 lastParent.Files[:idx],
 254                 lastParent.Files[idx+1:]...,
 255         )
 256         for _, parent := range parents {
 257                 parent.Size -= file.Size
 258         }
 259         return nil
 260 }
 261
 262 func (db *File) add(p string) error {
 263         cols := strings.Split(p, "/")
 264         cols, name := cols[:len(cols)-1], cols[len(cols)-1]
 265         var parent *File
 266         var err error
 267         if len(cols) != 0 {
 268                 parent, _, _, err = db.find(path.Join(cols...))
 269                 if err != nil {
 270                         return err
 271                 }
 272         } else {
 273                 parent = db
 274         }
 275         info, err := os.Stat(p)
 276         if err != nil {
 277                 return err
 278         }
 279         if info.IsDir() {
 280                 name += "/"
 281         }
 282         file := File{
 283                 Name:  name,
 284                 Size:  uint64(info.Size()),
 285                 Mtime: info.ModTime().Unix(),
 286         }
 287         parent.Files = append(parent.Files, file)
 288         sort.Sort(ByName(parent.Files))
 289         parent.Size += file.Size
 290         return nil
 291 }
 292
 293 func deoctalize(s string) string {
 294         chars := make([]byte, 0, len(s))
 295         for i := 0; i < len(s); i++ {
 296                 if s[i] == '\\' {
 297                         b, err := strconv.ParseUint("0"+s[i+1:i+1+3], 0, 8)
 298                         if err != nil {
 299                                 log.Fatalln(err)
 300                         }
 301                         chars = append(chars, byte(b))
 302                         i += 3
 303                 } else {
 304                         chars = append(chars, s[i])
 305                 }
 306         }
 307         return string(chars)
 308 }
 309
 310 func main() {
 311         dbPath := flag.String("db", ".glocate.db", "Path to state file (database)")
 312         doIndex := flag.Bool("index", false, "Initialize database")
 313         doUpdate := flag.Bool("update", false, "Update database by zfs-diff's output")
 314         showBeauty := flag.Bool("show-beauty", false, "Show beauty human-friendly listing")
 315         showSimple := flag.Bool("show-simple", false, "Show simple listing")
 316         stripPrefix := flag.String("strip-prefix", "", "Strip prefix from zfs-diff's output")
 317         flag.Parse()
 318         log.SetFlags(log.Ldate | log.Ltime | log.Lmicroseconds | log.Lshortfile)
 319
 320         if *doIndex {
 321                 files, size, err := walk(".")
 322                 if err != nil {
 323                         log.Fatalln(err)
 324                 }
 325                 db := File{Name: "./", Size: size, Files: files}
 326                 err = db.dump(*dbPath)
 327                 if err != nil {
 328                         log.Fatalln(err)
 329                 }
 330                 return
 331         }
 332
 333         db := load(*dbPath)
 334         if *doUpdate {
 335                 scanner := bufio.NewScanner(os.Stdin)
 336                 var t string
 337                 for scanner.Scan() {
 338                         t = scanner.Text()
 339                         if len(t) == 0 {
 340                                 continue
 341                         }
 342                         cols := strings.Split(t, "\t")
 343                         if len(cols) < 2 {
 344                                 log.Fatalln("bad zfs-diff format")
 345                         }
 346                         switch cols[0] {
 347                         case "-":
 348                                 name := deoctalize(strings.TrimPrefix(cols[1], *stripPrefix))
 349                                 if err := db.remove(name); err != nil {
 350                                         log.Println("can not -:", name, ":", err)
 351                                 }
 352                         case "+":
 353                                 name := deoctalize(strings.TrimPrefix(cols[1], *stripPrefix))
 354                                 if err := db.add(name); err != nil {
 355                                         log.Println("can not +:", name, ":", err)
 356                                 }
 357                         case "M":
 358                                 name := deoctalize(strings.TrimPrefix(cols[1], *stripPrefix))
 359                                 if name == "" {
 360                                         continue
 361                                 }
 362                                 file, _, _, err := db.find(name)
 363                                 if err != nil {
 364                                         log.Println("can not M:", name, ":", err)
 365                                         continue
 366                                 }
 367                                 info, err := os.Stat(name)
 368                                 if err != nil {
 369                                         log.Println("can not M:", name, ":", err)
 370                                         continue
 371                                 }
 372                                 if info.Mode().IsRegular() {
 373                                         file.Size = uint64(info.Size())
 374                                 }
 375                                 file.Mtime = info.ModTime().Unix()
 376                         case "R":
 377                                 if len(cols) != 3 {
 378                                         log.Fatalln("bad zfs-diff format for R")
 379                                 }
 380                                 name := deoctalize(strings.TrimPrefix(cols[1], *stripPrefix))
 381                                 if err := db.remove(name); err != nil {
 382                                         log.Println("can not R-:", name, ":", err)
 383                                         continue
 384                                 }
 385                                 name = deoctalize(strings.TrimPrefix(cols[2], *stripPrefix))
 386                                 if err := db.add(name); err != nil {
 387                                         log.Println("can not R+:", name, ":", err)
 388                                 }
 389                         default:
 390                                 log.Fatalln("bad zfs-diff format")
 391                         }
 392                 }
 393                 if err := scanner.Err(); err != nil {
 394                         log.Fatalln(err)
 395                 }
 396                 if err := db.dump(*dbPath); err != nil {
 397                         log.Fatalln(err)
 398                 }
 399                 return
 400         }
 401
 402         veryFirst := true
 403         if len(flag.Args()) > 0 {
 404                 root := flag.Args()[0]
 405                 if root[:2] == "./" {
 406                         root = root[2:]
 407                 }
 408                 if root[len(root)-1:] == "/" {
 409                         root = root[:len(root)-1]
 410                 }
 411                 file, _, _, err := db.find(root)
 412                 if err != nil {
 413                         log.Fatalln(err)
 414                 }
 415                 db = file
 416                 db.Name = root + "/"
 417                 veryFirst = false
 418         }
 419
 420         if *showBeauty {
 421                 db.listBeauty("", 0, false, veryFirst)
 422                 return
 423         }
 424         if *showSimple {
 425                 db.listSimple("", veryFirst)
 426                 return
 427         }
 428         db.listFiles("", veryFirst)
 429 }