Permalink
Browse files

SQL support; tested with SQLite

  • Loading branch information...
dsamarin committed May 11, 2018
1 parent 930e174 commit 5ef30b6d2550d9b0f83416709b5bcf42978451de
Showing with 185 additions and 87 deletions.
  1. +4 −17 .gitignore
  2. +27 −12 admin.go
  3. +13 −8 blacklist.go
  4. +1 −1 config.yml
  5. +118 −38 db.go
  6. +8 −3 main.go
  7. +14 −8 shot.go
@@ -1,20 +1,7 @@
# Binaries for programs and plugins
*.exe
*.dll
*.so
*.dylib
zerodrop
*.db
*.mmdb
*.sock
# Test binary, build with `go test -c`
*.test
# Output of the go coverage tool, specifically when used with LiteIDE
*.out
# Project-local glide cache, RE: https://github.com/Masterminds/glide/issues/736
.glide/
# GeoDB
GeoLite2-City.mmdb
# Uploads folder
uploads/
@@ -62,7 +62,7 @@ type AdminPageData struct {
Title string
LoggedIn bool
Config *ZerodropConfig
Entries []ZerodropEntry
Entries []*ZerodropEntry
}
// ServeLogin renders the login page.
@@ -166,10 +166,12 @@ func (a *AdminHandler) ServeNew(w http.ResponseWriter, r *http.Request) {
// Access information
entry.AccessExpire = r.FormValue("access_expire") != ""
entry.AccessExpireCount, _ = strconv.Atoi(r.FormValue("access_expire_count"))
entry.AccessBlacklist = ParseBlacklist(r.FormValue("blacklist"), a.Config.Databases)
entry.AccessBlacklist = ParseBlacklist(r.FormValue("blacklist"), a.Config.IPCat)
entry.AccessRedirectOnDeny = strings.TrimSpace(r.FormValue("access_redirect_on_deny"))
if err := a.DB.Create(entry); err == nil {
if err := a.DB.Update(entry); err != nil {
log.Printf("Error creating entry %s: %s", entry.Name, err)
} else {
log.Printf("Created entry %s", entry)
}
@@ -196,33 +198,46 @@ func (a *AdminHandler) ServeMain(w http.ResponseWriter, r *http.Request) {
case "train":
name := r.FormValue("name")
entry, ok := a.DB.Get(name)
if ok {
entry, err := a.DB.Get(name)
if err != nil {
log.Println(err)
} else {
entry.SetTraining(!entry.AccessTrain)
if err := a.DB.Update(entry); err != nil {
log.Println(err)
}
}
case "delete":
name := r.FormValue("name")
if name != "" {
a.DB.Remove(name)
log.Printf("Removed entry: %s", name)
err := a.DB.Remove(name)
if err != nil {
log.Println(err)
} else {
log.Printf("Removed entry: %s", name)
}
}
case "clear":
a.DB.Clear()
log.Println("Cleared all entries")
err := a.DB.Clear()
if err != nil {
log.Println(err)
} else {
log.Println("Cleared all entries")
}
}
http.Redirect(w, r, a.Config.Base+"admin/", 302)
return
}
data.Entries = a.DB.List()
var err error
data.Entries, err = a.DB.List()
interfaceTmpl := a.Templates.Lookup("admin-main.tmpl")
err := interfaceTmpl.ExecuteTemplate(w, "main", data)
if err != nil {
if interfaceTmpl.ExecuteTemplate(w, "main", data) != nil {
log.Println(err)
}
}
@@ -29,7 +29,7 @@ type BlacklistRule struct {
Network *net.IPNet
IP net.IP
Hostname string
Regexp *regexp.Regexp
Regexp string
Geofence *Geofence
Database string
}
@@ -71,9 +71,9 @@ func (i BlacklistRule) String() (value string) {
return
}
if i.Regexp != nil {
if i.Regexp != "" {
value += "~"
value += i.Regexp.String()
value += i.Regexp
if i.Comment != "" {
value += " # " + i.Comment
}
@@ -119,7 +119,7 @@ func (b Blacklist) String() string {
for index, item := range b.List {
if item.All || item.Network != nil || item.Geofence != nil ||
item.Database != "" || item.Hostname != "" || item.IP != nil ||
item.Regexp != nil {
item.Regexp != "" {
itemCount++
}
items[index+1] = item.String()
@@ -273,7 +273,7 @@ func ParseBlacklist(text string, dbconfig map[string]string) Blacklist {
case '~':
// An optional prefix "~" indicates a hostname regular expression match.
line = strings.TrimSpace(line[1:])
reg, err := regexp.Compile(line)
_, err := regexp.Compile(line)
if err != nil {
item.Comment = fmt.Sprintf(
"Error: %s: malformed regular expression: %s",
@@ -282,7 +282,7 @@ func ParseBlacklist(text string, dbconfig map[string]string) Blacklist {
continue
}
item.Regexp = reg
item.Regexp = line
blacklist.Add(item)
continue
}
@@ -359,13 +359,18 @@ func (b *Blacklist) Allow(ctx *BlacklistContext, ip net.IP) bool {
}
}
} else if item.Regexp != nil {
} else if item.Regexp != "" {
// Regular Expression
regex, err := regexp.Compile(item.Regexp)
if err != nil {
log.Printf("Error compiling regular expression: %s", err)
}
names, err := net.LookupAddr(ip.String())
if err != nil {
for _, name := range names {
name = strings.ToLower(name)
if item.Regexp.Match([]byte(name)) {
if regex.Match([]byte(name)) {
match = true
break
}
@@ -17,7 +17,7 @@ authdigest: 11a55ac5de2beb9146e01386dd978a13bb9b99388f5eb52e37f69a32e3d5f11e
geodb: GeoLite2-City.mmdb
# The databases to include for categorical blacklisting
databases:
ipcat:
datacenters: ../ipcat/datacenters.csv
tor: ../ipcat/tor.csv
156 db.go
@@ -1,15 +1,19 @@
package main
import (
"errors"
"sort"
"bytes"
"database/sql"
"encoding/gob"
"strconv"
"time"
_ "github.com/go-sql-driver/mysql"
_ "github.com/lib/pq"
_ "github.com/mattn/go-sqlite3"
)
// ZerodropEntry is a page entry.
type ZerodropEntry struct {
db *ZerodropDB
Name string // The request URI used to access this entry
URL string // The URL that this entry references
Filename string // The location of the file in the uploads directory
@@ -28,79 +32,155 @@ type ZerodropEntry struct {
// ZerodropDB represents a database connection.
// TODO: Use a persistent backend.
type ZerodropDB struct {
mapping map[string]ZerodropEntry
*sql.DB
GetStmt *sql.Stmt
ListStmt *sql.Stmt
CreateStmt *sql.Stmt
DeleteStmt *sql.Stmt
ClearStmt *sql.Stmt
}
// Connect opens a connection to the backend.
func (d *ZerodropDB) Connect() error {
d.mapping = map[string]ZerodropEntry{}
func (d *ZerodropDB) Connect(driver, source string) error {
db, err := sql.Open(driver, source)
if err != nil {
return err
}
gob.Register(&ZerodropEntry{})
if _, err := db.Exec(`CREATE TABLE IF NOT EXISTS entries (
name TEXT PRIMARY KEY NOT NULL,
creation INTEGER NOT NULL,
gob BLOB NOT NULL
)`); err != nil {
return err
}
d.GetStmt, err = db.Prepare(`SELECT gob FROM entries WHERE name = ?`)
if err != nil {
return err
}
d.ListStmt, err = db.Prepare(`SELECT gob FROM entries ORDER BY creation DESC`)
if err != nil {
return err
}
d.CreateStmt, err = db.Prepare(`REPLACE INTO entries (name, creation, gob) VALUES (?, ?, ?)`)
if err != nil {
return err
}
d.DeleteStmt, err = db.Prepare(`DELETE FROM entries WHERE name = ?`)
if err != nil {
return err
}
d.ClearStmt, err = db.Prepare(`DELETE FROM entries`)
if err != nil {
return err
}
d.DB = db
return nil
}
// Get returns the entry with the specified name.
func (d *ZerodropDB) Get(name string) (entry ZerodropEntry, ok bool) {
entry, ok = d.mapping[name]
return
func (d *ZerodropDB) Get(name string) (*ZerodropEntry, error) {
var data []byte
if err := d.GetStmt.QueryRow(name).Scan(&data); err != nil {
return nil, err
}
var entry *ZerodropEntry
dec := gob.NewDecoder(bytes.NewReader(data))
if err := dec.Decode(&entry); err != nil {
return nil, err
}
return entry, nil
}
// List returns a slice of all entries sorted by creation time,
// with the most recent first.
func (d *ZerodropDB) List() []ZerodropEntry {
list := []ZerodropEntry{}
func (d *ZerodropDB) List() ([]*ZerodropEntry, error) {
list := []*ZerodropEntry{}
rows, err := d.ListStmt.Query()
if err != nil {
return nil, err
}
defer rows.Close()
for rows.Next() {
var data []byte
if err := rows.Scan(&data); err != nil {
return nil, err
}
var entry *ZerodropEntry
dec := gob.NewDecoder(bytes.NewReader(data))
if err := dec.Decode(&entry); err != nil {
return nil, err
}
for _, entry := range d.mapping {
list = append(list, entry)
}
sort.Slice(list, func(i, j int) bool {
a := list[i].Creation
b := list[j].Creation
return a.After(b)
})
if err := rows.Err(); err != nil {
return nil, err
}
return list
return list, nil
}
// Create adds an entry to the database.
func (d *ZerodropDB) Create(entry *ZerodropEntry) error {
entry.db = d
d.mapping[entry.Name] = *entry
// Update adds an entry to the database.
func (d *ZerodropDB) Update(entry *ZerodropEntry) error {
var buffer bytes.Buffer
enc := gob.NewEncoder(&buffer)
if err := enc.Encode(entry); err != nil {
return err
}
if _, err := d.CreateStmt.Exec(entry.Name, entry.Creation.Unix(), buffer.Bytes()); err != nil {
return err
}
return nil
}
// Remove removes an entry from the database.
func (d *ZerodropDB) Remove(name string) {
delete(d.mapping, name)
func (d *ZerodropDB) Remove(name string) error {
if _, err := d.DeleteStmt.Exec(name); err != nil {
return err
}
return nil
}
// Clear resets the database by removing all entries.
func (d *ZerodropDB) Clear() {
d.Connect()
func (d *ZerodropDB) Clear() error {
if _, err := d.ClearStmt.Exec(); err != nil {
return err
}
return nil
}
// IsExpired returns true if the entry is expired
func (e *ZerodropEntry) IsExpired() bool {
return e.AccessExpire && (e.AccessCount >= e.AccessExpireCount)
}
// Update saves changes to the entry to the database it belongs to.
func (e *ZerodropEntry) Update() error {
if e.db != nil {
return e.db.Create(e)
}
return errors.New("No link to DB")
}
// SetTraining sets the AccessTrain flag
func (e *ZerodropEntry) SetTraining(train bool) error {
func (e *ZerodropEntry) SetTraining(train bool) {
e.AccessTrain = train
return e.Update()
}
// Access increases the access count for an entry.
func (e *ZerodropEntry) Access() error {
func (e *ZerodropEntry) Access() {
e.AccessCount++
return e.Update()
}
func (e *ZerodropEntry) String() string {
Oops, something went wrong.

0 comments on commit 5ef30b6

Please sign in to comment.