apiary/honeypot/store/postgres.go
2021-04-10 07:58:01 +02:00

132 lines
2.8 KiB
Go

package store
import (
"database/sql"
"fmt"
_ "github.com/jackc/pgx/v4/stdlib"
"github.uio.no/torjus/apiary/models"
)
const DBSchema = `
CREATE TABLE IF NOT EXISTS login_attempts(
id serial PRIMARY KEY,
date timestamptz,
remote_ip inet,
username text,
password text,
client_version text,
connection_uuid uuid,
country varchar(2)
);
`
type PostgresStore struct {
db *sql.DB
}
func NewPostgresStore(dsn string) (*PostgresStore, error) {
db, err := sql.Open("pgx", dsn)
if err != nil {
return nil, err
}
rs := &PostgresStore{
db: db,
}
return rs, nil
}
func (s *PostgresStore) InitDB() error {
_, err := s.db.Exec(DBSchema)
return err
}
func (s *PostgresStore) AddAttempt(l *models.LoginAttempt) error {
stmt := `INSERT INTO
login_attempts(date, remote_ip, username, password, client_version, connection_uuid, country)
VALUES ($1, $2, $3, $4, $5, $6, $7)
RETURNING id;`
tx, err := s.db.Begin()
if err != nil {
return err
}
defer tx.Rollback()
var id int
if err := tx.QueryRow(stmt, l.Date, l.RemoteIP.String(), l.Username, l.Password, l.SSHClientVersion, l.ConnectionUUID, l.Country).Scan(&id); err != nil {
return err
}
l.ID = id
return nil
}
func (s *PostgresStore) All() ([]models.LoginAttempt, error) {
stmt := `SELECT date, remote_ip, username, password, client_version, connection_uuid, country FROM login_attempts`
tx, err := s.db.Begin()
if err != nil {
return nil, err
}
rows, err := tx.Query(stmt)
if err != nil {
return nil, err
}
defer rows.Close()
var attempts []models.LoginAttempt
for rows.Next() {
var a models.LoginAttempt
if err := rows.Scan(&a.Date, &a.RemoteIP, &a.Username, &a.Password, &a.SSHClientVersion, &a.SSHClientVersion, &a.Country); err != nil {
return nil, err
}
attempts = append(attempts, a)
}
return attempts, nil
}
func (s *PostgresStore) Stats(statType LoginStats, limit int) ([]StatsResult, error) {
var stmt string
if statType == LoginStatsTotals {
return s.statsTotal(limit)
}
switch statType {
case LoginStatsCountry:
stmt = `select country, count(country) from login_attempts order by count desc`
case LoginStatsIP:
stmt = `select remote_ip, count(remote_ip) from login_attempts order by count desc`
case LoginStatsPasswords:
stmt = `select password, count(password) from login_attempts order by count desc`
case LoginStatsUsername:
stmt = `select username, count(username) from login_attempts order by count desc`
}
if limit > 0 {
stmt = fmt.Sprintf("%s limit %d", stmt, limit)
}
rows, err := s.db.Query(stmt)
if err != nil {
return nil, err
}
defer rows.Close()
var results []StatsResult
for rows.Next() {
var r StatsResult
if err := rows.Scan(&r.Name, &r.Count); err != nil {
return nil, err
}
results = append(results, r)
}
return results, nil
}
func (s *PostgresStore) statsTotal(limit int) ([]StatsResult, error) {
return nil, nil
}