git.r.bdr.sh - rbdr/mobius/blob - internal/mobius/ban.go

   1 package mobius
   2
   3 import (
   4         "gopkg.in/yaml.v3"
   5         "os"
   6         "path/filepath"
   7         "sync"
   8         "time"
   9 )
  10
  11 type BanFile struct {
  12         banList  map[string]*time.Time
  13         filePath string
  14
  15         sync.Mutex
  16 }
  17
  18 func NewBanFile(path string) (*BanFile, error) {
  19         bf := &BanFile{
  20                 filePath: path,
  21                 banList:  make(map[string]*time.Time),
  22         }
  23
  24         err := bf.Load()
  25
  26         return bf, err
  27 }
  28
  29 func (bf *BanFile) Load() error {
  30         bf.Lock()
  31         defer bf.Unlock()
  32
  33         bf.banList = make(map[string]*time.Time)
  34
  35         fh, err := os.Open(bf.filePath)
  36         if err != nil {
  37                 if os.IsNotExist(err) {
  38                         return nil
  39                 }
  40                 return err
  41         }
  42         defer fh.Close()
  43
  44         decoder := yaml.NewDecoder(fh)
  45         err = decoder.Decode(&bf.banList)
  46         if err != nil {
  47                 return err
  48         }
  49
  50         return nil
  51 }
  52
  53 func (bf *BanFile) Add(ip string, until *time.Time) error {
  54         bf.Lock()
  55         defer bf.Unlock()
  56
  57         bf.banList[ip] = until
  58
  59         out, err := yaml.Marshal(bf.banList)
  60         if err != nil {
  61                 return err
  62         }
  63
  64         return os.WriteFile(filepath.Join(bf.filePath), out, 0644)
  65 }
  66
  67 func (bf *BanFile) IsBanned(ip string) (bool, *time.Time) {
  68         bf.Lock()
  69         defer bf.Unlock()
  70
  71         if until, ok := bf.banList[ip]; ok {
  72                 return true, until
  73         }
  74
  75         return false, nil
  76 }