314 lines
9.1 KiB
Swift
314 lines
9.1 KiB
Swift
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<String>
|
|
|
|
var classifierVersion: Int = 0 {
|
|
didSet {
|
|
writeClassifierVersion()
|
|
}
|
|
}
|
|
|
|
/**
|
|
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
|
|
}
|
|
|
|
var capCount: Int {
|
|
caps.count
|
|
}
|
|
|
|
var imageCount: Int {
|
|
caps.reduce(0) { $0 + $1.value.count }
|
|
}
|
|
|
|
init(in folder: URL, writers: [String]) {
|
|
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)
|
|
}
|
|
|
|
func loadData() throws {
|
|
loadClassifierVersion(at: classifierVersionFile)
|
|
try loadCaps()
|
|
try updateCounts()
|
|
saveCapCountHTML()
|
|
}
|
|
|
|
private func loadClassifierVersion(at url: URL) {
|
|
guard fm.fileExists(atPath: url.path) else {
|
|
return
|
|
}
|
|
let content: String
|
|
do {
|
|
content = try String(contentsOf: url)
|
|
.trimmingCharacters(in: .whitespacesAndNewlines)
|
|
} catch {
|
|
log("Failed to read classifier version file: \(error)")
|
|
return
|
|
}
|
|
|
|
guard let value = Int(content) else {
|
|
log("Invalid classifier version: \(content)")
|
|
return
|
|
}
|
|
self.classifierVersion = value
|
|
}
|
|
|
|
private func writeClassifierVersion() {
|
|
do {
|
|
try "\(classifierVersion)".data(using: .utf8)!
|
|
.write(to: classifierVersionFile)
|
|
} catch {
|
|
log("Failed to save classifier version: \(error)")
|
|
}
|
|
}
|
|
|
|
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 =
|
|
"""
|
|
<body style="margin: 0;">
|
|
<div style="display: flex; justify-content: center;">
|
|
<div style="font-size: 60px; font-family: 'SF Pro Display',-apple-system,BlinkMacSystemFont,Helvetica,sans-serif; -webkit-font-smoothing: antialiased;">\(count)</div>
|
|
</div>
|
|
</body>
|
|
"""
|
|
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)")
|
|
}
|
|
}
|