git.oblomov.eu Git - rbot/blob - lib/rbot/maskdb.rb

   1 module Irc
   2   class NetmaskDb
   3     # helper backend class: generic nested radix tree
   4     class Tree
   5       attr_reader :pre, :chi
   6
   7       def initialize(pre = '', chi = Hash.new)
   8         @pre = pre
   9         @chi = chi
  10       end
  11
  12       def add(val, *prefs)
  13         str = prefs.shift or raise 'empty prefs'
  14         @pre = str.dup if @chi.empty?
  15
  16         n = 0
  17         @pre.size.times do
  18           break if @pre[n] != str[n]
  19           n += 1
  20         end
  21
  22         rest = str.slice(n .. -1)
  23
  24         if n != @pre.size
  25           prest = @pre.slice!(n .. -1)
  26           pc = prest.slice! 0
  27           @chi = {pc => Tree.new(prest, @chi)}
  28         end
  29
  30         c = rest.slice!(0)
  31
  32         if c
  33           (@chi[c] ||= Tree.new).add(val, rest, *prefs)
  34         else
  35           if prefs.empty?
  36             (@chi[''] ||= Array.new).push val
  37           else
  38             (@chi[''] ||= Tree.new).add(val, *prefs)
  39           end
  40         end
  41       end
  42
  43       def empty?
  44         @chi.empty?
  45       end
  46
  47       def remove(*prefs, &block)
  48         str = prefs.shift or raise 'empty prefs?'
  49         return nil unless @pre.empty? or str.index(@pre) == 0
  50         c = str.slice(@pre.size) || ''
  51         return nil unless @chi.include? c
  52         if c == ''
  53           if prefs.empty?
  54             @chi[c].reject!(&block)
  55           else
  56             @chi[c].remove(*prefs, &block)
  57           end
  58         else
  59           @chi[c].remove(str.slice((@pre.size + 1) .. -1), *prefs, &block)
  60         end
  61         @chi.delete(c) if @chi[c].empty?
  62
  63         if @chi.size == 1
  64           k = @chi.keys.shift
  65           return nil if k == ''
  66           @pre << k << @chi[k].pre
  67           @chi = @chi[k].chi
  68         end
  69       end
  70
  71       def find(*prefs)
  72         str = prefs.shift or raise 'empty prefs?'
  73         self.find_helper(str, *prefs) + self.find_helper(str.reverse, *prefs)
  74       end
  75
  76       protected
  77       def find_helper(*prefs)
  78         str = prefs.shift or raise 'empty prefs?'
  79         return [] unless @pre.empty? or str.index(@pre) == 0
  80         # puts "#{self.inspect}: #{str} == #{@pre} pfx matched"
  81         if !@chi.include? ''
  82           matches = []
  83         elsif Array === @chi['']
  84           matches = @chi['']
  85         else
  86           matches = @chi[''].find(*prefs)
  87         end
  88
  89         c = str.slice(@pre.size)
  90
  91         more = []
  92         if c and @chi.include?(c)
  93           more = @chi[c].find_helper(str.slice((@pre.size + 1) .. -1), *prefs)
  94         end
  95         return more + matches
  96       end
  97     end
  98
  99     # api wrapper for netmasks
 100
 101     def initialize
 102       @tree = Tree.new
 103     end
 104
 105     def cook_component(str)
 106       s = (str && !str.empty?) ? str : '*'
 107       l = s.index(/[\?\*]/)
 108       if l
 109         l2 = s.size - s.rindex(/[\?\*]/) - 1
 110         if l2 > l
 111           s = s.reverse
 112           l = l2
 113         end
 114
 115         return (l > 0) ? s.slice(0 .. (l - 1)) : ''
 116       else
 117         return s
 118       end
 119     end
 120
 121     def mask2keys(m)
 122       md = m.downcased
 123       [md.host, md.user, md.nick].map { |c| cook_component(c) }
 124     end
 125
 126     def add(user, *masks)
 127       masks.each do |m|
 128         debug "adding user #{user} with mask #{m.fullform}"
 129         @tree.add([user, m], *mask2keys(m))
 130       end
 131     end
 132
 133     def remove(user, mask)
 134       debug "trying to remove user #{user} with mask #{mask}"
 135       @tree.remove(*mask2keys(mask)) do |val|
 136         val[0] == user and val[1].fullform == mask.fullform
 137       end
 138     end
 139
 140     def metric(iu, bu, mask)
 141       ret = nil
 142       if iu.matches? mask
 143         ret = iu.fullform.length - mask.fullform.length
 144         ret += 10 if bu.transient?
 145       end
 146       return ret
 147     end
 148
 149     def find(iu)
 150       debug "find(#{iu.fullform})"
 151       iud = iu.downcased
 152       matches = @tree.find(iud.host, iud.user, iud.nick).uniq.map do |val|
 153         m = metric(iu, *val)
 154         m ? [val[0], m] : nil
 155       end.compact.sort { |a, b| a[1] <=> a[1] }
 156       debug "matches: " + (matches.map do |m|
 157         "#{m[0].username}: [#{m[1]}]"
 158       end.join(', '))
 159       return matches.empty? ? nil : matches[0][0]
 160     end
 161   end
 162 end