go
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 
 

271 lines
7.0 KiB

  1. package main
  2. import (
  3. "database/sql"
  4. "flag"
  5. "fmt"
  6. "io/ioutil"
  7. "log"
  8. "net/http"
  9. "os"
  10. "strconv"
  11. "time"
  12. "github.com/gorilla/mux"
  13. _ "github.com/mattn/go-sqlite3"
  14. "gopkg.in/yaml.v2"
  15. "honnef.co/go/netdb"
  16. )
  17. const (
  18. CleanRecreate = "Recreate" // completely remove DB file and create again
  19. CleanFull = "Full"
  20. CleanLastDay = "LastDay"
  21. CleanLastWeek = "LastWeek"
  22. CleanLastMonth = "LastMonth"
  23. CleanNever = "Never"
  24. )
  25. type CleanType string
  26. type AppConfig struct {
  27. Ident string `yaml:"ident,omitempty"` // ident of this instance, if empty autogenerate every run
  28. DB struct {
  29. MasterDB string `yaml:"masterDB"` // master sqlite database
  30. Clean CleanType `yaml:"clean,omitempty"` // cleanup at startup
  31. } `yaml:"DB"`
  32. API struct {
  33. TokenWrite string `yaml:"tokenWrite,omitempty"` // auth for write requests, allow anonymous if empty
  34. TokenRead string `yaml:"tokenRead,omitempty"` // auth for read requests, allow anonymous if empty
  35. Listen string `yaml:"listen,omitempty"` // open host:port, default 0.0.0.0:21745
  36. } `yaml:"API"`
  37. }
  38. type IP struct {
  39. ID string
  40. Updated string
  41. Origin string
  42. Protocol string
  43. Port int
  44. IP string
  45. BlockTable string
  46. }
  47. var (
  48. configFilename = flag.String("config", "", "Configuration filename")
  49. blockIP = BlockIPdb{}
  50. )
  51. type BlockIPdb struct {
  52. HandlerDB *sql.DB
  53. }
  54. func (d *BlockIPdb) dbLoadFile(f string) error {
  55. log.Println("Loading master DB file", f)
  56. db, err := sql.Open("sqlite3", f)
  57. if err != nil {
  58. return err
  59. }
  60. d.HandlerDB = db
  61. err = d.dbInit()
  62. if err != nil {
  63. return err
  64. }
  65. return nil
  66. }
  67. func (d *BlockIPdb) dbInit() error {
  68. checkTable := "SELECT name FROM sqlite_master WHERE type='table' AND name='ip';"
  69. var checkTableResult string
  70. err := d.HandlerDB.QueryRow(checkTable).Scan(&checkTableResult)
  71. if err != nil {
  72. if err == sql.ErrNoRows {
  73. log.Println("Initializing database")
  74. createTable := "create table ip (id integer not null primary key, origin text, updated int, protocol text, port int, ip text unique, blocktable text);"
  75. _, err = d.HandlerDB.Exec(createTable)
  76. if err != nil {
  77. return err
  78. }
  79. } else {
  80. return err
  81. }
  82. }
  83. return nil
  84. }
  85. func main() {
  86. flag.Parse()
  87. if *configFilename == "" {
  88. log.Fatalln("Set configuration filename")
  89. }
  90. // read settings from file
  91. log.Println("Loading config file", *configFilename)
  92. appConfig := AppConfig{}
  93. yamlFile, err := ioutil.ReadFile(*configFilename)
  94. if err != nil {
  95. log.Fatalf("Config read error: %v\n", err)
  96. }
  97. err = yaml.Unmarshal(yamlFile, &appConfig)
  98. if err != nil {
  99. log.Fatalf("Config format error: %v\n", err)
  100. }
  101. err = blockIP.dbLoadFile(appConfig.DB.MasterDB)
  102. defer blockIP.HandlerDB.Close()
  103. if err != nil {
  104. log.Fatalf("Database error: %v\n", err)
  105. }
  106. switch appConfig.DB.Clean {
  107. case CleanRecreate:
  108. log.Println("Recreate DB option set")
  109. blockIP.HandlerDB.Close()
  110. os.Remove(appConfig.DB.MasterDB)
  111. err = blockIP.dbLoadFile(appConfig.DB.MasterDB)
  112. defer blockIP.HandlerDB.Close()
  113. if err != nil {
  114. log.Fatalf("Database error: %v\n", err)
  115. }
  116. case CleanFull:
  117. log.Println("Full cleanup option set")
  118. _, err = blockIP.HandlerDB.Exec("drop table ip;")
  119. if err != nil {
  120. log.Fatalf("Database error: %v\n", err)
  121. }
  122. err = blockIP.dbInit()
  123. if err != nil {
  124. log.Fatalf("Database error: %v\n", err)
  125. }
  126. case CleanLastDay:
  127. log.Println("Save only last day cleanup option set")
  128. _, err = blockIP.HandlerDB.Exec("delete from ip WHERE updated <= (strftime('%s','now') - (60 * ?));", 1440)
  129. if err != nil {
  130. log.Fatalf("Database error: %v\n", err)
  131. }
  132. case CleanLastWeek:
  133. log.Println("Save only last week option set")
  134. _, err = blockIP.HandlerDB.Exec("delete from ip WHERE updated <= (strftime('%s','now') - (60 * ?));", 1440*7)
  135. if err != nil {
  136. log.Fatalf("Database error: %v\n", err)
  137. }
  138. case CleanLastMonth:
  139. log.Println("Save only last 30 days option set")
  140. _, err = blockIP.HandlerDB.Exec("delete from ip WHERE updated <= (strftime('%s','now') - (60 * ?));", 1440*30)
  141. if err != nil {
  142. log.Fatalf("Database error: %v\n", err)
  143. }
  144. }
  145. _, err = blockIP.HandlerDB.Exec("VACUUM;")
  146. if err != nil {
  147. log.Fatalf("Database error: %v\n", err)
  148. }
  149. queryRecordsNum := "SELECT count(*) FROM ip;"
  150. var queryRecordsNumResult int
  151. err = blockIP.HandlerDB.QueryRow(queryRecordsNum).Scan(&queryRecordsNumResult)
  152. if err != nil {
  153. log.Fatalf("Query error: %v\n", err)
  154. }
  155. log.Println("IP records in database:", queryRecordsNumResult)
  156. r := mux.NewRouter()
  157. r.HandleFunc("/", HandlerHealth)
  158. r.HandleFunc("/set/{origin}/{blocktable}/{ip}/{proto}/{port}", HandlerSet)
  159. r.HandleFunc("/list/{blocktable}/{minutes}", HandlerList)
  160. // http.Handle("/", r)
  161. srv := &http.Server{
  162. Handler: r,
  163. Addr: appConfig.API.Listen,
  164. WriteTimeout: 5 * time.Second,
  165. ReadTimeout: 5 * time.Second,
  166. }
  167. log.Fatal(srv.ListenAndServe())
  168. }
  169. func HandlerHealth(w http.ResponseWriter, r *http.Request) {
  170. w.WriteHeader(http.StatusOK)
  171. }
  172. func HandlerList(w http.ResponseWriter, r *http.Request) {
  173. vars := mux.Vars(r)
  174. blockTable := vars["blocktable"]
  175. minutes, err := strconv.Atoi(vars["minutes"])
  176. if err != nil {
  177. log.Println(err)
  178. return
  179. }
  180. log.Println("Listing minutes:", minutes)
  181. rows := &sql.Rows{}
  182. if minutes != 0 {
  183. rows, err = blockIP.HandlerDB.Query("select ip from ip WHERE blocktable=? and updated >= (strftime('%s','now') - (60 * ?));", blockTable, minutes)
  184. } else {
  185. rows, err = blockIP.HandlerDB.Query("select ip from ip where blocktable=?", blockTable)
  186. }
  187. defer rows.Close()
  188. w.WriteHeader(http.StatusOK)
  189. for rows.Next() {
  190. var ip string
  191. err = rows.Scan(&ip)
  192. if err != nil {
  193. log.Println(err)
  194. }
  195. fmt.Fprintf(w, "%v\n", ip)
  196. }
  197. err = rows.Err()
  198. if err != nil {
  199. log.Println(err)
  200. }
  201. }
  202. func HandlerSet(w http.ResponseWriter, r *http.Request) {
  203. vars := mux.Vars(r)
  204. ip := IP{
  205. IP: vars["ip"],
  206. Protocol: vars["proto"],
  207. BlockTable: vars["blocktable"],
  208. Origin: vars["origin"],
  209. }
  210. port, err := strconv.Atoi(vars["port"])
  211. if err != nil {
  212. serv := netdb.GetServByName(vars["port"], netdb.GetProtoByName(ip.Protocol))
  213. if serv != nil {
  214. ip.Port = serv.Port
  215. } else {
  216. log.Println(err)
  217. return
  218. }
  219. } else {
  220. ip.Port = port
  221. }
  222. // query ip before insert
  223. var count int
  224. err = blockIP.HandlerDB.QueryRow("select count(ip) from ip WHERE blocktable=? and ip=?;", ip.BlockTable, ip.IP).Scan(&count)
  225. if err != nil {
  226. log.Println(err)
  227. }
  228. if count > 0 {
  229. log.Println("State: UPDATE", "IP:", ip.IP, "Origin:", ip.Origin, "Port:", ip.Port, "Protocol:", ip.Protocol)
  230. _, err := blockIP.HandlerDB.Exec("update ip set origin=?, updated=strftime('%s','now') where blocktable=? and ip=?", ip.Origin, ip.BlockTable, ip.IP)
  231. if err != nil {
  232. log.Println(err)
  233. }
  234. } else {
  235. log.Println("State: NEW", "IP:", ip.IP, "Origin:", ip.Origin, "Port:", ip.Port, "Protocol:", ip.Protocol)
  236. _, err := blockIP.HandlerDB.Exec("insert into ip(origin, ip, protocol, port, updated, blocktable) values(?, ?, ?, ?, strftime('%s','now'), ?)", ip.Origin, ip.IP, ip.Protocol, ip.Port, ip.BlockTable)
  237. if err != nil {
  238. log.Println(err)
  239. }
  240. }
  241. w.WriteHeader(http.StatusOK)
  242. }