A More Secure Internet Connection for Your Home https://fen.gg
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.
 
 
 
 
 
 

200 lines
5.1 KiB

  1. package model
  2. //
  3. // Fengg Security Gateway Server Application
  4. // Copyright (C) 2020 Lukas Matt <support@fen.gg>
  5. //
  6. // This program is free software: you can redistribute it and/or modify
  7. // it under the terms of the GNU General Public License as published by
  8. // the Free Software Foundation, either version 3 of the License.
  9. //
  10. // This program is distributed in the hope that it will be useful,
  11. // but WITHOUT ANY WARRANTY; without even the implied warranty of
  12. // MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
  13. // GNU General Public License for more details.
  14. //
  15. // You should have received a copy of the GNU General Public License
  16. // along with this program. If not, see <https://www.gnu.org/licenses/>.
  17. //
  18. import (
  19. "fmt"
  20. "time"
  21. "github.com/jmoiron/sqlx"
  22. _ "github.com/lib/pq"
  23. )
  24. const (
  25. PatternCategoryQueryTmpl = `
  26. SELECT * FROM pattern_categories %s;`
  27. PatternCategoryJoinedQueryTmpl = `
  28. SELECT %s
  29. FROM pattern_categories as category
  30. INNER JOIN pattern_category_lists as list ON category.id=list.category_id
  31. INNER JOIN patterns as pattern ON pattern.id=list.pattern_id
  32. %s;`
  33. PatternCategoryNamedUpdateTmpl = `
  34. UPDATE pattern_categories
  35. SET updated_at=now(), enabled=:enabled, count=:count, parent_id=:parent_id
  36. %s;`
  37. PatternCategoryNamedInsertTmpl = `
  38. INSERT INTO pattern_categories (
  39. updated_at, name, parent_id, enabled, count
  40. ) VALUES (
  41. now(), :name, :parent_id, :enabled, :count
  42. ) RETURNING id;`
  43. PatternCategoryCopyCSVQuery = `
  44. copy pattern_categories
  45. from '/usr/share/fengg-pattern/pattern-categories.csv'
  46. delimiter ',' csv header;`
  47. PatternCategoryNamedFixCounterTmpl = `
  48. with category_count as (
  49. select count(*) as count from pattern_categories as category
  50. inner join pattern_category_lists as list on category.id=list.category_id
  51. inner join patterns as pattern on pattern.id=list.pattern_id
  52. where category.name = :name
  53. )
  54. update pattern_categories
  55. set count = category_count.count
  56. from category_count
  57. where name = :name;`
  58. )
  59. type PatternCategory struct {
  60. ID uint `db:"id" json:"id"`
  61. UpdatedAt time.Time `db:"updated_at" json:"updatedAt"`
  62. Name string `db:"name" json:"name"`
  63. ParentID *uint `db:"parent_id" json:"-"`
  64. Enabled bool `db:"enabled" json:"enabled"`
  65. Count int `db:"count" json:"count"`
  66. }
  67. type PatternCategories []*PatternCategory
  68. func NewPatternCategory() *PatternCategory {
  69. return &PatternCategory{
  70. UpdatedAt: time.Now(),
  71. }
  72. }
  73. func (categories *PatternCategories) FindAll() error {
  74. db, err := sqlx.Connect(dbDriver, dbConnect)
  75. if err != nil {
  76. return err
  77. }
  78. defer db.Close()
  79. return db.Select(categories, fmt.Sprintf(PatternCategoryQueryTmpl, ""))
  80. }
  81. // FixCategoryCounter will update the counter column. This happens normally after an update.
  82. func (categories PatternCategories) FixCategoryCounter() error {
  83. err := categories.FindAll()
  84. if err != nil {
  85. return err
  86. }
  87. db, err := sqlx.Connect(dbDriver, dbConnect)
  88. if err != nil {
  89. return err
  90. }
  91. defer db.Close()
  92. tx := db.MustBegin()
  93. for _, category := range categories {
  94. _, err = tx.NamedExec(PatternCategoryNamedFixCounterTmpl, *category)
  95. if err != nil {
  96. return err
  97. }
  98. }
  99. return tx.Commit()
  100. }
  101. func (categories *PatternCategories) FindByEnabled(enabled bool) error {
  102. db, err := sqlx.Connect(dbDriver, dbConnect)
  103. if err != nil {
  104. return err
  105. }
  106. defer db.Close()
  107. return db.Select(categories, fmt.Sprintf(
  108. PatternCategoryQueryTmpl, "WHERE enabled = $1"), enabled)
  109. }
  110. func (categories PatternCategories) FindPatterns() (patterns Patterns, err error) {
  111. for _, category := range categories {
  112. result, err := category.FindPatterns()
  113. if err != nil {
  114. return patterns, err
  115. }
  116. patterns = append(patterns, result...)
  117. }
  118. return
  119. }
  120. func (category PatternCategory) FindPatterns() (patterns Patterns, err error) {
  121. db, err := sqlx.Connect(dbDriver, dbConnect)
  122. if err != nil {
  123. return nil, err
  124. }
  125. defer db.Close()
  126. return patterns, db.Select(&patterns, fmt.Sprintf(
  127. PatternCategoryJoinedQueryTmpl,
  128. "pattern.*",
  129. "WHERE category.id = $1",
  130. ), category.ID)
  131. }
  132. func (category *PatternCategory) Exists() bool {
  133. db, err := sqlx.Connect(dbDriver, dbConnect)
  134. if err != nil {
  135. logger.Error().Err(err).Msg("cannot connect to database")
  136. return false
  137. }
  138. defer db.Close()
  139. return db.Get(category,
  140. fmt.Sprintf(PatternCategoryQueryTmpl, "WHERE id = $1 OR name = $2"),
  141. category.ID, category.Name) == nil
  142. }
  143. func (category *PatternCategory) Update() error {
  144. db, err := sqlx.Connect(dbDriver, dbConnect)
  145. if err != nil {
  146. return err
  147. }
  148. defer db.Close()
  149. _, err = db.NamedExec(fmt.Sprintf(
  150. PatternCategoryNamedUpdateTmpl, "WHERE id=:id"), category)
  151. return err
  152. }
  153. func (category *PatternCategory) Create() error {
  154. db, err := sqlx.Connect(dbDriver, dbConnect)
  155. if err != nil {
  156. return err
  157. }
  158. defer db.Close()
  159. rows, err := db.NamedQuery(PatternCategoryNamedInsertTmpl, category)
  160. if err != nil {
  161. return err
  162. }
  163. defer rows.Close()
  164. for rows.Next() {
  165. err = rows.Scan(&category.ID)
  166. if err != nil {
  167. return err
  168. }
  169. }
  170. return nil
  171. }