messages.go 2.2 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182
  1. package data
  2. import (
  3. "context"
  4. "database/sql"
  5. "github.com/owncast/owncast/db"
  6. "github.com/owncast/owncast/models"
  7. log "github.com/sirupsen/logrus"
  8. )
  9. // GetMessagesCount will return the number of messages in the database.
  10. func GetMessagesCount() int64 {
  11. query := `SELECT COUNT(*) FROM messages`
  12. rows, err := _db.Query(query)
  13. if err != nil || rows.Err() != nil {
  14. return 0
  15. }
  16. defer rows.Close()
  17. var count int64
  18. for rows.Next() {
  19. if err := rows.Scan(&count); err != nil {
  20. return 0
  21. }
  22. }
  23. return count
  24. }
  25. // CreateBanIPTable will create the IP ban table if needed.
  26. func CreateBanIPTable(db *sql.DB) {
  27. createTableSQL := ` CREATE TABLE IF NOT EXISTS ip_bans (
  28. "ip_address" TEXT NOT NULL PRIMARY KEY,
  29. "notes" TEXT,
  30. "created_at" TIMESTAMP DEFAULT CURRENT_TIMESTAMP
  31. );`
  32. stmt, err := db.Prepare(createTableSQL)
  33. if err != nil {
  34. log.Fatal("error creating ip ban table", err)
  35. }
  36. defer stmt.Close()
  37. if _, err := stmt.Exec(); err != nil {
  38. log.Fatal("error creating ip ban table", err)
  39. }
  40. }
  41. // BanIPAddress will persist a new IP address ban to the datastore.
  42. func BanIPAddress(address, note string) error {
  43. return _datastore.GetQueries().BanIPAddress(context.Background(), db.BanIPAddressParams{
  44. IpAddress: address,
  45. Notes: sql.NullString{String: note, Valid: true},
  46. })
  47. }
  48. // IsIPAddressBanned will return if an IP address has been previously blocked.
  49. func IsIPAddressBanned(address string) (bool, error) {
  50. blocked, error := _datastore.GetQueries().IsIPAddressBlocked(context.Background(), address)
  51. return blocked > 0, error
  52. }
  53. // GetIPAddressBans will return all the banned IP addresses.
  54. func GetIPAddressBans() ([]models.IPAddress, error) {
  55. result, err := _datastore.GetQueries().GetIPAddressBans(context.Background())
  56. if err != nil {
  57. return nil, err
  58. }
  59. response := []models.IPAddress{}
  60. for _, ip := range result {
  61. response = append(response, models.IPAddress{
  62. IPAddress: ip.IpAddress,
  63. Notes: ip.Notes.String,
  64. CreatedAt: ip.CreatedAt.Time,
  65. })
  66. }
  67. return response, err
  68. }
  69. // RemoveIPAddressBan will remove a previously banned IP address.
  70. func RemoveIPAddressBan(address string) error {
  71. return _datastore.GetQueries().RemoveIPAddressBan(context.Background(), address)
  72. }