Skip to content

Commit a891111

Browse files
committed
#248 Fix double encoded metrics in dblist during migration
1 parent 27c04fb commit a891111

File tree

1 file changed

+41
-29
lines changed

1 file changed

+41
-29
lines changed

thehive-backend/app/models/Migration.scala

+41-29
Original file line numberDiff line numberDiff line change
@@ -145,6 +145,7 @@ class Migration(
145145
def unapply(data: String): Option[Array[Byte]] = Try(java.util.Base64.getDecoder.decode(data)).toOption
146146
}
147147

148+
// store attachment id and check to prevent document already exists error
148149
var dataIds = Set.empty[String]
149150
def containsOrAdd(id: String) = {
150151
dataIds.synchronized {
@@ -155,45 +156,56 @@ class Migration(
155156
}
156157
}
157158
}
159+
158160
val mainHasher = Hasher(mainHash)
159161
val extraHashers = Hasher(mainHash +: extraHashes: _*)
160162
Seq(
163+
// store alert attachment in datastore
161164
Operation((f: String Source[JsObject, NotUsed]) {
162-
case "alert" f("alert").flatMapConcat { alert
163-
val artifactsAndData = Future.traverse((alert \ "artifacts").asOpt[List[JsObject]].getOrElse(Nil)) { artifact
164-
(artifact \ "data").asOpt[String]
165-
.collect {
166-
case AlertSrv.dataExtractor(filename, contentType, data @ Base64(rawData))
167-
val attachmentId = mainHasher.fromByteArray(rawData).head.toString()
168-
ds.getEntity(datastoreName, s"${attachmentId}_0")
169-
.map(_ Nil)
170-
.recover {
171-
case _ if containsOrAdd(attachmentId) Nil
172-
case _
173-
Seq(Json.obj(
174-
"_type" datastoreName,
175-
"_id" s"${attachmentId}_0",
176-
"data" data))
177-
}
178-
.map { dataEntity
179-
val attachment = Attachment(filename, extraHashers.fromByteArray(rawData), rawData.length.toLong, contentType, attachmentId)
180-
(artifact - "data" + ("attachment" Json.toJson(attachment))) dataEntity
181-
}
165+
case "alert" f("alert").flatMapConcat { alert
166+
val artifactsAndData = Future.traverse((alert \ "artifacts").asOpt[List[JsObject]].getOrElse(Nil)) { artifact
167+
(artifact \ "data").asOpt[String]
168+
.collect {
169+
case AlertSrv.dataExtractor(filename, contentType, data @ Base64(rawData))
170+
val attachmentId = mainHasher.fromByteArray(rawData).head.toString()
171+
ds.getEntity(datastoreName, s"${attachmentId}_0")
172+
.map(_ Nil)
173+
.recover {
174+
case _ if containsOrAdd(attachmentId) Nil
175+
case _
176+
Seq(Json.obj(
177+
"_type" datastoreName,
178+
"_id" s"${attachmentId}_0",
179+
"data" data))
180+
}
181+
.map { dataEntity
182+
val attachment = Attachment(filename, extraHashers.fromByteArray(rawData), rawData.length.toLong, contentType, attachmentId)
183+
(artifact - "data" + ("attachment" Json.toJson(attachment))) dataEntity
184+
}
185+
}
186+
.getOrElse(Future.successful(artifact Nil))
187+
}
188+
Source.fromFuture(artifactsAndData)
189+
.mapConcat { ad
190+
val updatedAlert = alert + ("artifacts" JsArray(ad.map(_._1)))
191+
updatedAlert :: ad.flatMap(_._2)
182192
}
183-
.getOrElse(Future.successful(artifact Nil))
184193
}
185-
Source.fromFuture(artifactsAndData)
186-
.mapConcat { ad
187-
val updatedAlert = alert + ("artifacts" JsArray(ad.map(_._1)))
188-
updatedAlert :: ad.flatMap(_._2)
189-
}
190-
}
191-
case other f(other)
192-
}),
194+
case other f(other)
195+
}),
196+
// Fix alert status
193197
mapAttribute("alert", "status") {
194198
case JsString("Update") JsString("Updated")
195199
case JsString("Ignore") JsString("Ignored")
196200
case other other
201+
},
202+
// Fix double encode of metrics
203+
mapEntity("dblist") {
204+
case dblist if (dblist \ "dblist").asOpt[String].contains("case_metrics")
205+
(dblist \ "value").asOpt[String].map(Json.parse).fold(dblist) { value
206+
dblist + ("value" value)
207+
}
208+
case other other
197209
})
198210
}
199211

0 commit comments

Comments
 (0)