import Foundation import Vapor final class CapServer { // MARK: Paths private let imageFolder: URL /// The file where the database of caps is stored private let dbFile: URL /// The file to store the HTML info of the cap count private let htmlFile: URL private let classifierVersionFile: URL private let classifierFile: URL private let fm = FileManager.default // MARK: Caps private var writers: Set var classifierVersion: Int { get { do { let content = try String(contentsOf: classifierVersionFile) .trimmingCharacters(in: .whitespacesAndNewlines) guard let value = Int(content) else { log("Invalid classifier version: \(content)") return 0 } return value } catch { log("Failed to read classifier version file: \(error)") return 0 } } set { do { try "\(newValue)".data(using: .utf8)! .write(to: classifierVersionFile) } catch { log("Failed to save classifier version: \(error)") } } } /** The time to wait for changes to be written to disk. This delay is used to prevent file writes for each small update to the caps. */ private let saveDelay: TimeInterval = 1 /** The time when a save should occur. No save is necessary if this property is `nil`. */ private var nextSaveTime: Date? private var caps = [Int: Cap]() { didSet { scheduleSave() } } var nextClassifierVersion: Int { caps.values.compactMap { $0.classifierVersion }.max() ?? 1 } init(in folder: URL, writers: [String]) throws { self.imageFolder = folder.appendingPathComponent("images") self.dbFile = folder.appendingPathComponent("caps.json") self.htmlFile = folder.appendingPathComponent("count.html") self.classifierVersionFile = folder.appendingPathComponent("classifier.version") self.classifierFile = folder.appendingPathComponent("classifier.mlmodel") self.writers = Set(writers) var isDirectory: ObjCBool = false guard fm.fileExists(atPath: folder.path, isDirectory: &isDirectory), isDirectory.boolValue else { log("Public directory \(folder.path) is not a folder, or doesn't exist") throw CapError.invalidConfiguration } try loadCaps() try updateCounts() saveCapCountHTML() } private func loadCaps() throws { do { let data = try Data(contentsOf: dbFile) caps = try JSONDecoder().decode([Cap].self, from: data) .reduce(into: [:]) { $0[$1.id] = $1 } } catch { log("Failed to load caps: \(error)") throw error } log("\(caps.count) caps loaded") } private func scheduleSave() { nextSaveTime = Date().addingTimeInterval(saveDelay) DispatchQueue.global().asyncAfter(deadline: .now() + saveDelay) { self.performScheduledSave() } } private func performScheduledSave() { guard let date = nextSaveTime else { // No save necessary, or already saved return } guard date < Date() else { // Save pushed to future return } do { try saveCaps() nextSaveTime = nil } catch { // Attempt save again scheduleSave() } } private func saveCaps() throws { let data = try JSONEncoder().encode(caps.values.sorted()) try data.write(to: dbFile) } private func saveCapCountHTML() { let count = caps.count let content = """

\(count)

""" try? content.data(using: .utf8)!.write(to: htmlFile) } // MARK: Paths func folder(of cap: Int) -> URL { imageFolder.appendingPathComponent(String(format: "%04d", cap)) } func file(of cap: Int, version: Int) -> URL { folder(of: cap).appendingPathComponent(String(format: "%04d-%02d.jpg", cap, version)) } // MARK: Authentication func hasAuthorization(for key: String) -> Bool { // Note: This is not a constant-time compare, so there may be an opportunity // for timing attack here. Sets perform hashed lookups, so this may be less of an issue, // and we're not doing anything critical in this application. // Worst case, an unauthorized person with a lot of free time and energy to hack this system // is able to change contents of the database, which are backed up in any case. writers.contains(key) } // MARK: Counts private func updateCounts() throws { do { caps = try caps.mapValues { var cap = $0 cap.count = try count(of: $0.id) return cap } } catch { log("Failed to update counts: \(error)") throw error } } func countImages(in folder: URL) throws -> Int { try fm.contentsOfDirectory(at: folder, includingPropertiesForKeys: nil) .filter({ $0.pathExtension == "jpg" }).count } /** Get the image count of a cap. */ func count(of cap: Int) throws -> Int { let f = folder(of: cap) guard fm.fileExists(atPath: f.path) else { return 0 } return try countImages(in: f) } // MARK: Images /** Save a cap image to disk. Automatically creates the image name with the current image count. - Parameter data: The image data - Parameter cap: The id of the cap. - Throws: `CapError.unknownId`, if the cap doesn't exist. `CapError.dataInconsistency` if an image already exists for the current count. */ func save(image data: Data, for cap: Int) throws { guard caps[cap] != nil else { throw CapError.unknownId } var id = 0 let capFolder = folder(of: cap) var f = file(of: cap, version: id) if fm.fileExists(atPath: capFolder.path) { while fm.fileExists(atPath: f.path) { id += 1 f = file(of: cap, version: id) } } else { try fm.createDirectory(at: capFolder, withIntermediateDirectories: true) } try data.write(to: f) caps[cap]!.count = try count(of: cap) log("Added image \(id) for cap \(cap)") } func switchMainImage(to version: Int, for cap: Int) throws { let file2 = file(of: cap, version: version) guard fm.fileExists(atPath: file2.path) else { log("No image \(version) for cap \(cap)") throw CapError.invalidFile } caps[cap]?.mainImage = version log("Switched cap \(cap) to version \(version)") } func addOrUpdate(_ cap: Cap) throws { if let existingCap = caps[cap.id] { update(existingCap, with: cap) } else { try add(cap) } } private func add(_ cap: Cap) throws { guard cap.mainImage == 0 else { throw CapError.invalidData } var cap = cap cap.count = 0 cap.classifierVersion = nextClassifierVersion caps[cap.id] = cap saveCapCountHTML() log("Added cap \(cap.id) '\(cap.name)'") } private func update(_ existingCap: Cap, with cap: Cap) { var updatedCap = existingCap if cap.name != "" { updatedCap.name = cap.name } let url = file(of: existingCap.id, version: cap.mainImage) if fm.fileExists(atPath: url.path) { updatedCap.mainImage = cap.mainImage } if let color = cap.color { updatedCap.color = color } caps[existingCap.id] = updatedCap log("Updated cap \(existingCap.id)") } func updateTrainedClasses(content: String) { let trainedCaps = content .components(separatedBy: "\n") .compactMap(Int.init) let version = classifierVersion for cap in trainedCaps { if caps[cap]?.classifierVersion == nil { caps[cap]?.classifierVersion = version } } log("Updated \(trainedCaps.count) classifier classes") } func save(classifier: Data, version: Int) throws { do { try classifier.write(to: classifierFile) } catch { log("Failed to write classifier: \(error)") throw Abort(.internalServerError) } classifierVersion = version log("Updated classifier to version \(version)") } }