add a Site model class

SOBotics · Oct 18, 2017 · 5d4ade4 · 5d4ade4
1 parent f27ba71
commit 5d4ade4
Show file tree

Hide file tree

Showing 17 changed files with 342 additions and 79 deletions.
diff --git a/Sources/FireAlarm/Commands/Filters/CommandAddSite.swift b/Sources/FireAlarm/Commands/Filters/CommandAddSite.swift
@@ -30,17 +30,17 @@ class CommandAddSite: Command {
             return
         }
 
-        guard let site: Int = try reporter.staticDB.run(
+        guard let site = try reporter.staticDB.run(
             "SELECT id FROM sites WHERE domain = ? OR apiSiteParameter = ?",
             siteName, siteName
-            ).first?.column(at: 0) else {
+            ).first.map(Site.from) else {
 
                 reply("That does not look like a site on which I run.")
                 return
         }
 
-        message.room.thresholds[site] = threshold
+        message.room.thresholds[site.id] = threshold
 
-        reply("Added \(siteName) to this room's sites with threshold \(threshold).")
+        reply("Added \(site.domain) to this room's sites with threshold \(threshold).")
     }
 }
diff --git a/Sources/FireAlarm/Commands/Filters/CommandCheckSites.swift b/Sources/FireAlarm/Commands/Filters/CommandCheckSites.swift
@@ -22,9 +22,7 @@ class CommandCheckSites: Command {
         }
 
         let siteNames: [String] = try message.room.thresholds.keys.map {
-            try reporter.staticDB.run(
-                "SELECT domain FROM sites WHERE id = ?", $0
-                ).first?.column(at: 0) ?? "<unknown site \($0)>"
+            try Site.with(id: $0, db: reporter.staticDB)?.domain ?? "<unknown site \($0)>"
         }
 
         reply("This room reports posts from \(formatArray(siteNames, conjunction: "and")).")

diff --git a/Sources/FireAlarm/Commands/Filters/CommandCheckThreshold.swift b/Sources/FireAlarm/Commands/Filters/CommandCheckThreshold.swift
@@ -26,9 +26,7 @@ class CommandCheckThreshold: Command {
             )
         } else {
             let siteNames: [String] = try message.room.thresholds.keys.map {
-                try reporter.staticDB.run(
-                    "SELECT domain FROM sites WHERE id = ?", $0
-                    ).first?.column(at: 0) ?? "<unknown site \($0)>"
+                try Site.with(id: $0, db: reporter.staticDB)?.domain ?? "<unknown site \($0)>"
             }
             let siteThresholds = Array(message.room.thresholds.values.map(String.init))
 

diff --git a/Sources/FireAlarm/Commands/Filters/CommandRemoveSite.swift b/Sources/FireAlarm/Commands/Filters/CommandRemoveSite.swift
@@ -21,18 +21,18 @@ class CommandRemoveSite: Command {
     override func run() throws {
         guard
             let siteName = arguments.first,
-            let site: Int = try reporter.staticDB.run(
+            let site = try reporter.staticDB.run(
                 "SELECT id FROM sites WHERE domain = ? OR apiSiteParameter = ?",
                 siteName, siteName
-                ).first?.column(at: 0),
-            message.room.thresholds[site] != nil else {
+                ).first.map(Site.from),
+            message.room.thresholds[site.id] != nil else {
 
                 reply("Please enter a site and a threshold.")
                 return
         }
 
-        message.room.thresholds[site] = nil
+        message.room.thresholds[site.id] = nil
 
-        reply("Removed \(siteName) from this room's sites.")
+        reply("Removed \(site.domain) from this room's sites.")
     }
 }
diff --git a/Sources/FireAlarm/Commands/Filters/CommandSetThreshold.swift b/Sources/FireAlarm/Commands/Filters/CommandSetThreshold.swift
@@ -38,20 +38,21 @@ class CommandSetThreshold: Command {
                 message.reply("Which site would you like to change threshold for?")
             }
         } else {
-            guard let siteID = try reporter.staticDB.run("SELECT id FROM sites " +
+            guard let site = try reporter.staticDB.run("SELECT id FROM sites " +
                 "WHERE apiSiteParameter = ? OR domain = ?",
                                       arguments.last!, arguments.last!
-                ).first?.column(at: 0) as Int? else {
+                ).first.map(Site.from) else {
 
                     message.reply("That does not look like a site on which I run.")
                     return
             }
-            guard message.room.thresholds[siteID] != nil else {
+            guard message.room.thresholds[site.id] != nil else {
                 message.reply("I do not report posts from that site here.")
                 return
             }
 
-            message.room.thresholds[siteID] = newThreshold
+            message.room.thresholds[site.id] = newThreshold
+            reply("Set threshold for `\(site.domain)` to \(newThreshold).")
         }
     }
 }
diff --git a/Sources/FireAlarm/Commands/Reports/CommandCheckPost.swift b/Sources/FireAlarm/Commands/Reports/CommandCheckPost.swift
@@ -32,13 +32,9 @@ class CommandCheckPost: Command {
         }
 
         guard
-            let site = try reporter.staticDB.run(
-                "SELECT id FROM sites WHERE domain = ?",
-                siteDomain
-                ).first?.column(at: 0) as Int? else {
-
-                    reply("That does not look like a site on which I run.")
-                    return
+            let site = try Site.with(domain: siteDomain, db: reporter.staticDB) else {
+                reply("That does not look like a site on which I run.")
+                return
         }
 
         guard

diff --git a/Sources/FireAlarm/Commands/Reports/CommandReport.swift b/Sources/FireAlarm/Commands/Reports/CommandReport.swift
@@ -38,11 +38,7 @@ class CommandReport: Command {
         }
 
         guard
-            let site = try reporter.staticDB.run(
-                "SELECT id FROM sites WHERE domain = ?",
-                siteDomain
-                ).first?.column(at: 0) as Int? else {
-
+            let site = try Site.with(domain: siteDomain, db: reporter.staticDB) else {
                     reply("That does not look like a site on which I run.")
                     return
         }

diff --git a/Sources/FireAlarm/Filter.swift b/Sources/FireAlarm/Filter.swift
@@ -43,6 +43,6 @@ struct FilterResult {
 
 protocol Filter: class {
     init(reporter: Reporter)
-    func check(_ post: Question, site: Int) throws -> FilterResult?
+    func check(_ post: Question, site: Site) throws -> FilterResult?
 	func save() throws
 }
diff --git a/Sources/FireAlarm/Filters/FilterBlacklistedKeywords.swift b/Sources/FireAlarm/Filters/FilterBlacklistedKeywords.swift
@@ -38,15 +38,15 @@ class FilterBlacklistedKeywords: Filter {
             if FileManager.default.fileExists(atPath: keywordURL.path) {
                 print("Backing up blacklisted_users.json.")
                 do {
-                    try FileManager.default.moveItem(at: keywordURL, to: saveDirURL.appendingPathComponent("blacklisted_users.json.bak"))
+                    try FileManager.default.moveItem(at: keywordURL, to: saveDirURL.appendingPathComponent("blacklisted_keywords.json.bak"))
                 } catch {
                     handleError(error, "while backing up the blacklisted keywords")
                 }
             }
         }
     }
 
-    func check(_ post: Question, site: Int) -> FilterResult? {
+    func check(_ post: Question, site: Site) -> FilterResult? {
         guard let content = post.body else {
             print("No body for \(post.id.map { String($0) } ?? "<no ID>")!")
             return nil

diff --git a/Sources/FireAlarm/Filters/FilterBlacklistedUsernames.swift b/Sources/FireAlarm/Filters/FilterBlacklistedUsernames.swift
@@ -46,7 +46,7 @@ class FilterBlacklistedUsernames: Filter {
         }
     }
 
-    func check(_ post: Question, site: Int) -> FilterResult? {
+    func check(_ post: Question, site: Site) -> FilterResult? {
         guard let name = post.owner?.display_name else {
             print("No username for \(post.id.map { String($0) } ?? "<no ID>")!")
             return nil

diff --git a/Sources/FireAlarm/Filters/FilterMisleadingLinks.swift b/Sources/FireAlarm/Filters/FilterMisleadingLinks.swift
@@ -14,7 +14,7 @@ import Dispatch
 class FilterMisleadingLinks: Filter {
     required init (reporter: Reporter) {}
 
-    func check(_ post: Question, site: Int) -> FilterResult? {
+    func check(_ post: Question, site: Site) -> FilterResult? {
         do {
             let regex = try NSRegularExpression(pattern:
                 "<a href=\"([^\"]*)\" rel=\"nofollow(?: noreferrer)?\">\\s*([^<\\s]*)(?=\\s*</a>)", options: []

diff --git a/Sources/FireAlarm/Filters/FilterNaiveBayes.swift b/Sources/FireAlarm/Filters/FilterNaiveBayes.swift
@@ -35,16 +35,8 @@ class FilterNaiveBayes: Filter {
         print ("Loading Naive Bayes filter...")
     }
 
-    func check(_ post: Question, site: Int) throws -> FilterResult? {
-        guard let initialProbability = try reporter.staticDB.run(
-            "SELECT initialProbability FROM sites WHERE id = ?", site
-            ).first?.column(at: 0) as Double? else {
-
-                print("No initialProbability for site \(site)")
-                return nil
-        }
-
-        var trueProbability = Double(initialProbability)
+    func check(_ post: Question, site: Site) throws -> FilterResult? {
+        var trueProbability = Double(site.initialProbability)
         var falseProbability = Double(1 - trueProbability)
         var postWords = [String]()
         var checkedWords = [String]()
@@ -78,7 +70,7 @@ class FilterNaiveBayes: Filter {
                 "SELECT * FROM words " +
                 "WHERE site = ? " +
                 "AND word = ?;",
-                site, postWord
+                site.id, postWord
                 ).first.map(Word.init) else {
 
                 continue

diff --git a/Sources/FireAlarm/Models/Site.swift b/Sources/FireAlarm/Models/Site.swift
@@ -8,7 +8,7 @@
 import Foundation
 import SwiftChatSE
 
-struct Site: Codable {
+struct Site: Codable, Hashable {
     var id: Int
     var apiSiteParameter: String
     var domain: String
@@ -18,7 +18,27 @@ struct Site: Codable {
         return try RowDecoder().decode(Site.self, from: row)
     }
 
+    ///Helper function used by methods such as with(id:db:).
+    ///- warning: The `name` parameter is vulnerable to SQL injection; make sure it's valid!
+    static private func with<T: DatabaseType>(parameter: T, name: String, db: DatabaseConnection) throws -> Site? {
+        return try db.run("SELECT * FROM sites WHERE \(name) = ?;", parameter).first.map(Site.from)
+    }
+
+    ///Returns the Site with the specified ID.
     static func with(id: Int, db: DatabaseConnection) throws -> Site? {
-        return try db.run("SELECT * FROM sites WHERE id = ?;", id).first.map(Site.from)
+        return try with(parameter: id, name: "id", db: db)
+    }
+
+    ///Returns the Site with the specified domain.
+    static func with(domain: String, db: DatabaseConnection) throws -> Site? {
+        return try with(parameter: domain, name: "domain", db: db)
     }
+
+    ///Returns the Site with the specified apiSiteParameter.
+    static func with(apiSiteParameter: String, db: DatabaseConnection) throws -> Site? {
+        return try with(parameter: apiSiteParameter, name: "apiSiteParameter", db: db)
+    }
+
+    var hashValue: Int { return id.hashValue }
+    static func ==(lhs: Site, rhs: Site) -> Bool { return lhs.id == rhs.id }
 }
diff --git a/Sources/FireAlarm/PostFetcher.swift b/Sources/FireAlarm/PostFetcher.swift
@@ -12,7 +12,7 @@ import SwiftStack
 import Dispatch
 
 class PostFetcher {
-    var postsToCheck = [(id: Int, site: Int)]()
+    var postsToCheck = [(id: Int, site: Site)]()
 
     var queue = DispatchQueue(label: "Filter", attributes: [.concurrent])
 
@@ -72,7 +72,7 @@ class PostFetcher {
                         !posts.contains { $0.id == post.id && $0.site == post.site }
                     }
 
-                    var postsBySite = [Int:[Int]]()
+                    var postsBySite = [Site:[Int]]()
                     for post in posts {
                         if postsBySite[post.site] != nil {
                             postsBySite[post.site]!.append(post.id)
@@ -82,17 +82,9 @@ class PostFetcher {
                     }
 
                     for (site, posts) in postsBySite {
-                        guard let apiSiteParameter = try self.staticDB.run(
-                            "SELECT apiSiteParameter FROM sites WHERE id = ?",
-                            site
-                            ).first?.column(at: 0) as String? else {
-
-                                throw QuestionProcessingError.siteLookupFailed(siteID: site)
-                        }
-
                         for post in try apiClient.fetchQuestions(
                             posts,
-                            parameters: ["site":apiSiteParameter]
+                            parameters: ["site":site.apiSiteParameter]
                             ).items ?? [] {
 
                                 //don't report posts that are more than a day old
@@ -214,23 +206,19 @@ class PostFetcher {
                     throw QuestionProcessingError.noDataObject(json: string)
                 }
 
-                guard let site = data["apiSiteParameter"] as? String else {
+                guard let apiSiteParameter = data["apiSiteParameter"] as? String else {
                     throw QuestionProcessingError.noSite(json: string)
                 }
 
-                guard let siteID = try staticDB.run(
-                    "SELECT id FROM sites WHERE apiSiteParameter = ?",
-                    site
-                    ).first?.column(at: 0) as Int? else {
+                guard let site = try Site.with(apiSiteParameter: apiSiteParameter, db: staticDB) else {
                         return
                 }
 
                 guard let id = data["id"] as? Int else {
                     throw QuestionProcessingError.noQuestionID(json: string)
                 }
 
-
-                postsToCheck.append((id: id, site: siteID))
+                postsToCheck.append((id: id, site: site))
                 //print("Another post has been recieved.  There are now \(postsToCheck.count) posts to check.")
 
             } catch {

diff --git a/Sources/FireAlarm/Reporter.swift b/Sources/FireAlarm/Reporter.swift
@@ -151,11 +151,11 @@ class Reporter {
         postFetcher = PostFetcher(rooms: rooms, reporter: self, staticDB: staticDB)
     }
 
-    func checkPost(_ post: Question, site: Int) throws -> [FilterResult] {
+    func checkPost(_ post: Question, site: Site) throws -> [FilterResult] {
         return try filters.flatMap { try $0.check(post, site: site) }
     }
 
-    @discardableResult func checkAndReportPost(_ post: Question, site: Int) throws -> ReportResult {
+    @discardableResult func checkAndReportPost(_ post: Question, site: Site) throws -> ReportResult {
         let results = try checkPost(post, site: site)
 
         return try report(post: post, site: site, reasons: results)
@@ -190,10 +190,10 @@ class Reporter {
     }
 
     ///Reports a post if it has not been recently reported.  Returns either .reported or .alreadyReported.
-    func report(post: Question, site: Int, reasons: [FilterResult]) throws -> ReportResult {
+    func report(post: Question, site: Site, reasons: [FilterResult]) throws -> ReportResult {
         var status: ReportResult.Status = .notBad
 
-        try queue.sync {
+        queue.sync {
             guard let id = post.id else {
                 print("No post ID!")
                 status = .notBad
@@ -242,7 +242,7 @@ class Reporter {
                 let reasons = reasons.filter {
                     if case .bayesianFilter(let difference) = $0.type {
                         bayesianDifference = difference
-                        return difference < room.thresholds[site] ?? Int.min
+                        return difference < room.thresholds[site.id] ?? Int.min
                     }
                     return true
                 }
@@ -253,15 +253,9 @@ class Reporter {
 
                 reported = true
 
-                guard let domain = try staticDB.run("SELECT domain FROM sites WHERE id = ?", site)
-                    .first?.column(at: 0) as String? else {
-
-                        throw ReportError.missingSite(id: site)
-                }
-
                 let header = reasons.map { $0.header }.joined(separator: ", ")
                 let message = "[ [\(botName)](\(stackAppsLink)) ] " +
-                    "[tag:\(tags.first ?? "tagless")] \(header) [\(title)](//\(domain)/q/\(id)) " +
+                    "[tag:\(tags.first ?? "tagless")] \(header) [\(title)](//\(site.domain)/q/\(id)) " +
                     room.notificationString(tags: tags, reasons: reasons)
 
                 room.postMessage(message, completion: {message in