Skip to content

Commit

Permalink
#1410 Add extra data for cases and observables
Browse files Browse the repository at this point in the history
  • Loading branch information
To-om committed Jun 28, 2020
1 parent 19a1452 commit e12a821
Show file tree
Hide file tree
Showing 11 changed files with 198 additions and 83 deletions.
5 changes: 3 additions & 2 deletions dto/src/main/scala/org/thp/thehive/dto/v1/Case.scala
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@ package org.thp.thehive.dto.v1
import java.util.Date

import org.thp.scalligraph.controllers.WithParser
import play.api.libs.json.{Json, OFormat, OWrites}
import play.api.libs.json.{JsObject, Json, OFormat, OWrites}

case class InputCase(
title: String,
Expand Down Expand Up @@ -46,7 +46,8 @@ case class OutputCase(
status: String,
summary: Option[String] = None,
assignee: Option[String],
customFields: Set[OutputCustomFieldValue] = Set.empty
customFields: Set[OutputCustomFieldValue] = Set.empty,
extraData: JsObject
)

object OutputCase {
Expand Down
2 changes: 1 addition & 1 deletion dto/src/main/scala/org/thp/thehive/dto/v1/Observable.scala
Original file line number Diff line number Diff line change
Expand Up @@ -47,7 +47,7 @@ case class OutputObservable(
ioc: Boolean,
sighted: Boolean,
message: Option[String],
stats: JsObject
extraData: JsObject
)

object OutputObservable {
Expand Down
4 changes: 2 additions & 2 deletions thehive/app/org/thp/thehive/controllers/v0/Conversion.scala
Original file line number Diff line number Diff line change
Expand Up @@ -139,7 +139,7 @@ object Conversion {
.withFieldConst(_._type, "case")
.withFieldComputed(_.id, _._id)
.withFieldRenamed(_.number, _.caseId)
.withFieldRenamed(_.user, _.owner)
.withFieldRenamed(_.assignee, _.owner)
.withFieldRenamed(_._updatedAt, _.updatedAt)
.withFieldRenamed(_._updatedBy, _.updatedBy)
.withFieldRenamed(_._createdAt, _.createdAt)
Expand Down Expand Up @@ -192,7 +192,7 @@ object Conversion {
.withFieldConst(_._type, "case")
.withFieldComputed(_.id, _._id)
.withFieldRenamed(_.number, _.caseId)
.withFieldRenamed(_.user, _.owner)
.withFieldRenamed(_.assignee, _.owner)
.withFieldRenamed(_._updatedAt, _.updatedAt)
.withFieldRenamed(_._updatedBy, _.updatedBy)
.withFieldRenamed(_._createdAt, _.createdAt)
Expand Down
7 changes: 3 additions & 4 deletions thehive/app/org/thp/thehive/controllers/v1/CaseCtrl.scala
Original file line number Diff line number Diff line change
Expand Up @@ -42,10 +42,9 @@ class CaseCtrl @Inject() (
"page",
FieldsParser[OutputParam], {
case (OutputParam(from, to, extraData), caseSteps, authContext) =>
caseSteps
.richPage(from, to, extraData.contains("total")) { c =>
c.richCaseWithCustomRenderer(caseStatsRenderer(extraData - "total")(authContext, db, caseSteps.graph))(authContext)
}
caseSteps.richPage(from, to, extraData.contains("total")) {
_.richCaseWithCustomRenderer(caseStatsRenderer(extraData - "total")(authContext, db, caseSteps.graph))(authContext)
}
}
)
override val outputQuery: Query = Query.outputWithContext[RichCase, CaseSteps]((caseSteps, authContext) => caseSteps.richCase(authContext))
Expand Down
87 changes: 87 additions & 0 deletions thehive/app/org/thp/thehive/controllers/v1/CaseRenderer.scala
Original file line number Diff line number Diff line change
@@ -0,0 +1,87 @@
package org.thp.thehive.controllers.v1

import java.util.{Map => JMap}

import gremlin.scala.{__, By, Graph, GremlinScala, Key, Vertex}
import org.thp.scalligraph.auth.AuthContext
import org.thp.scalligraph.models.Database
import org.thp.scalligraph.steps.StepsOps._
import org.thp.scalligraph.steps.Traversal
import org.thp.thehive.models.AlertCase
import org.thp.thehive.services.CaseSteps
import play.api.libs.json._

import scala.collection.JavaConverters._

trait CaseRenderer {

def observableStats(
caseSteps: CaseSteps
)(implicit authContext: AuthContext): Traversal[JsValue, JsValue] =
caseSteps
.share
.observables
.count
.map(count => Json.obj("count" -> count))

def taskStats(caseSteps: CaseSteps)(implicit authContext: AuthContext): Traversal[JsValue, JsValue] =
caseSteps
.share
.tasks
.active
.groupCount(By(Key[String]("status")))
.map { statusAgg =>
val (total, result) = statusAgg.asScala.foldLeft(0L -> JsObject.empty) {
case ((t, r), (k, v)) => (t + v) -> (r + (k -> JsNumber(v.toInt)))
}
result + ("total" -> JsNumber(total))
}

def alertStats(caseSteps: CaseSteps): Traversal[JsValue, JsValue] =
caseSteps
.inTo[AlertCase]
.group(By(Key[String]("type")), By(Key[String]("source")))
.map { alertAgg =>
JsArray(
alertAgg
.asScala
.flatMap {
case (tpe, listOfSource) =>
listOfSource.asScala.map(s => Json.obj("type" -> tpe, "source" -> s))
}
.toSeq
)
}

def isOwnerStats(
caseSteps: CaseSteps
)(implicit authContext: AuthContext): Traversal[JsValue, JsValue] =
caseSteps.origin.name.map(org => JsBoolean(org == authContext.organisation))

def shareCountStats(caseSteps: CaseSteps): Traversal[JsValue, JsValue] =
caseSteps.organisations.count.map(JsNumber.apply(_))

def caseStatsRenderer(extraData: Set[String])(
implicit authContext: AuthContext,
db: Database,
graph: Graph
): CaseSteps => Traversal[JsObject, JsObject] = {
def addData(f: CaseSteps => Traversal[JsValue, JsValue]): GremlinScala[JMap[String, JsValue]] => GremlinScala[JMap[String, JsValue]] =
_.by(f(new CaseSteps(__[Vertex])).raw.traversal)

if (extraData.isEmpty) _.constant(JsObject.empty)
else {
val dataName = extraData.toSeq
dataName
.foldLeft[CaseSteps => GremlinScala[JMap[String, JsValue]]](_.raw.project(dataName.head, dataName.tail: _*)) {
case (f, "observableStats") => f.andThen(addData(observableStats))
case (f, "taskStats") => f.andThen(addData(taskStats))
case (f, "alerts") => f.andThen(addData(alertStats))
case (f, "isOwner") => f.andThen(addData(isOwnerStats))
case (f, "shareCount") => f.andThen(addData(shareCountStats))
case (f, _) => f.andThen(_.by(__.constant(JsNull).traversal))
}
.andThen(f => Traversal(f.map(m => JsObject(m.asScala))))
}
}
}
40 changes: 37 additions & 3 deletions thehive/app/org/thp/thehive/controllers/v1/Conversion.scala
Original file line number Diff line number Diff line change
Expand Up @@ -68,12 +68,22 @@ object Conversion {
.withFieldComputed(_.customFields, _.customFields.map(_.toOutput).toSet)
.withFieldComputed(_.tags, _.tags.map(_.toString).toSet)
.withFieldComputed(_.status, _.status.toString)
.withFieldRenamed(_.user, _.assignee)
.withFieldConst(_.extraData, JsObject.empty)
.transform
)

implicit val caseWithStatsOutput: Renderer[(RichCase, JsObject)] =
Renderer.json[(RichCase, JsObject), OutputCase](_._1.toOutput) // TODO add stats
Renderer.json[(RichCase, JsObject), OutputCase] { caseWithExtraData =>
caseWithExtraData
._1
.into[OutputCase]
.withFieldConst(_._type, "Case")
.withFieldComputed(_.customFields, _.customFields.map(_.toOutput).toSet)
.withFieldComputed(_.tags, _.tags.map(_.toString).toSet)
.withFieldComputed(_.status, _.status.toString)
.withFieldConst(_.extraData, caseWithExtraData._2)
.transform
}

implicit class InputCaseOps(inputCase: InputCase) {

Expand Down Expand Up @@ -252,6 +262,15 @@ object Conversion {
.transform
)

implicit class InputObservableOps(inputObservable: InputObservable) {
def toObservable: Observable =
inputObservable
.into[Observable]
.withFieldComputed(_.ioc, _.ioc.getOrElse(false))
.withFieldComputed(_.sighted, _.sighted.getOrElse(false))
.withFieldComputed(_.tlp, _.tlp.getOrElse(2))
.transform
}
implicit val observableOutput: Renderer.Aux[RichObservable, OutputObservable] = Renderer.json[RichObservable, OutputObservable](richObservable =>
richObservable
.into[OutputObservable]
Expand All @@ -266,10 +285,25 @@ object Conversion {
.withFieldComputed(_.tags, _.tags.map(_.toString).toSet)
.withFieldComputed(_.data, _.data.map(_.data))
.withFieldComputed(_.attachment, _.attachment.map(_.toOutput))
.withFieldConst(_.stats, JsObject.empty)
.withFieldConst(_.extraData, JsObject.empty)
.transform
)

implicit val observableWithExtraData: Renderer.Aux[(RichObservable, JsObject), OutputObservable] =
Renderer.json[(RichObservable, JsObject), OutputObservable] {
case (richObservable, extraData) =>
richObservable
.into[OutputObservable]
.withFieldConst(_._type, "case_artifact")
.withFieldComputed(_.dataType, _.`type`.name)
.withFieldComputed(_.startDate, _.observable._createdAt)
.withFieldComputed(_.tags, _.tags.map(_.toString).toSet)
.withFieldComputed(_.data, _.data.map(_.data))
.withFieldComputed(_.attachment, _.attachment.map(_.toOutput))
.withFieldConst(_.extraData, extraData)
.transform
}

implicit val logOutput: Renderer.Aux[RichLog, OutputLog] = Renderer.json[RichLog, OutputLog](richLog =>
richLog
.into[OutputLog]
Expand Down
30 changes: 6 additions & 24 deletions thehive/app/org/thp/thehive/controllers/v1/ObservableCtrl.scala
Original file line number Diff line number Diff line change
Expand Up @@ -7,15 +7,13 @@ import org.thp.scalligraph.models.Database
import org.thp.scalligraph.query.{ParamQuery, PropertyUpdater, PublicProperty, Query}
import org.thp.scalligraph.steps.PagedResult
import org.thp.scalligraph.steps.StepsOps._
import org.thp.thehive.controllers.v0.Conversion._
import org.thp.thehive.dto.v0.InputObservable
import org.thp.thehive.models._
import org.thp.thehive.services._
import play.api.Logger
import play.api.libs.json.JsObject
import play.api.mvc.{Action, AnyContent, Results}

import scala.util.Success
import org.thp.thehive.controllers.v1.Conversion._
import org.thp.thehive.dto.v1.InputObservable

@Singleton
class ObservableCtrl @Inject() (
Expand All @@ -42,14 +40,10 @@ class ObservableCtrl @Inject() (
override val pageQuery: ParamQuery[OutputParam] = Query.withParam[OutputParam, ObservableSteps, PagedResult[(RichObservable, JsObject)]](
"page",
FieldsParser[OutputParam], {
case (OutputParam(from, to, withStats), observableSteps, authContext) =>
observableSteps
.richPage(from, to, withTotal = true) {
case o if withStats =>
o.richObservableWithCustomRenderer(observableStatsRenderer(authContext, db, observableSteps.graph))
case o =>
o.richObservable.map(_ -> JsObject.empty)
}
case (OutputParam(from, to, extraData), observableSteps, authContext) =>
observableSteps.richPage(from, to, extraData.contains("total")) {
_.richObservableWithCustomRenderer(observableStatsRenderer(extraData - "total")(authContext, db, observableSteps.graph))
}
}
)
override val outputQuery: Query = Query.output[RichObservable, ObservableSteps](_.richObservable)
Expand Down Expand Up @@ -112,18 +106,6 @@ class ObservableCtrl @Inject() (
.map(_ => Results.NoContent)
}

def findSimilar(obsId: String): Action[AnyContent] =
entryPoint("find similar")
.authRoTransaction(db) { _ => implicit graph =>
val observables = observableSrv
.getByIds(obsId)
.similar
.richObservableWithCustomRenderer(observableLinkRenderer)
.toList

Success(Results.Ok(observables.toJson))
}

def bulkUpdate: Action[AnyContent] =
entryPoint("bulk update")
.extract("input", FieldsParser.update("observable", publicProperties))
Expand Down
67 changes: 40 additions & 27 deletions thehive/app/org/thp/thehive/controllers/v1/ObservableRenderer.scala
Original file line number Diff line number Diff line change
@@ -1,50 +1,63 @@
package org.thp.thehive.controllers.v1

import gremlin.scala.{__, By, Graph, Key, Vertex}
import javax.inject.Named
import java.util.{Map => JMap}

import gremlin.scala.{__, By, Graph, GremlinScala, Key, Vertex}
import org.thp.scalligraph.auth.AuthContext
import org.thp.scalligraph.models.Database
import org.thp.scalligraph.steps.StepsOps._
import org.thp.scalligraph.steps.Traversal
import org.thp.thehive.controllers.v0.Conversion._
import org.thp.thehive.services.ObservableSteps
import play.api.libs.json.{JsObject, Json}
import play.api.libs.json._

import scala.collection.JavaConverters._

trait ObservableRenderer {

def observableStatsRenderer(
implicit authContext: AuthContext,
@Named("with-thehive-schema") db: Database,
graph: Graph
): ObservableSteps => Traversal[JsObject, JsObject] =
_.project(
_.apply(By(new ObservableSteps(__[Vertex]).shares.organisation.name.fold.raw))
.and(
By(
new ObservableSteps(__[Vertex])
.similar
.visible
.groupCount(By(Key[Boolean]("ioc")))
.raw
)
)
).map {
case (organisationNames, stats) =>
def seenStats(observableSteps: ObservableSteps)(implicit authContext: AuthContext): Traversal[JsValue, JsValue] =
observableSteps
.similar
.visible
.groupCount(By(Key[Boolean]("ioc")))
.map { stats =>
val m = stats.asScala
val nTrue = m.get(true).fold(0L)(_.toLong)
val nFalse = m.get(false).fold(0L)(_.toLong)
Json.obj(
"shares" -> organisationNames.asScala,
"seen" -> (nTrue + nFalse),
"ioc" -> (nTrue > 0)
"seen" -> (nTrue + nFalse),
"ioc" -> (nTrue > 0)
)
}
}

def sharesStats(observableSteps: ObservableSteps): Traversal[JsValue, JsValue] =
observableSteps.shares.organisation.name.fold.map(orgs => Json.toJson(orgs.asScala))

def observableLinkRenderer: ObservableSteps => Traversal[JsObject, JsObject] =
_.coalesce(
def observableLinks(observableSteps: ObservableSteps): Traversal[JsValue, JsValue] =
observableSteps.coalesce(
_.alert.richAlert.map(a => Json.obj("alert" -> a.toJson)),
_.`case`.richCaseWithoutPerms.map(c => Json.obj("case" -> c.toJson))
)

def observableStatsRenderer(extraData: Set[String])(
implicit authContext: AuthContext,
db: Database,
graph: Graph
): ObservableSteps => Traversal[JsObject, JsObject] = {
def addData(f: ObservableSteps => Traversal[JsValue, JsValue]): GremlinScala[JMap[String, JsValue]] => GremlinScala[JMap[String, JsValue]] =
_.by(f(new ObservableSteps(__[Vertex])).raw.traversal)

if (extraData.isEmpty) _.constant(JsObject.empty)
else {
val dataName = extraData.toSeq
dataName
.foldLeft[ObservableSteps => GremlinScala[JMap[String, JsValue]]](_.raw.project(dataName.head, dataName.tail: _*)) {
case (f, "seen") => f.andThen(addData(seenStats))
case (f, "shares") => f.andThen(addData(sharesStats))
case (f, "links") => f.andThen(addData(observableLinks))
case (f, _) => f.andThen(_.by(__.constant(JsNull).traversal))
}
.andThen(f => Traversal(f.map(m => JsObject(m.asScala))))
}
}
}
2 changes: 1 addition & 1 deletion thehive/app/org/thp/thehive/models/Case.scala
Original file line number Diff line number Diff line change
Expand Up @@ -100,7 +100,7 @@ case class RichCase(
tags: Seq[Tag with Entity],
impactStatus: Option[String],
resolutionStatus: Option[String],
user: Option[String],
assignee: Option[String],
customFields: Seq[RichCustomField],
userPermissions: Set[Permission]
) {
Expand Down
2 changes: 1 addition & 1 deletion thehive/test/org/thp/thehive/DatabaseBuilder.scala
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@ import org.thp.scalligraph.RichOption
import org.thp.scalligraph.auth.AuthContext
import org.thp.scalligraph.controllers._
import org.thp.scalligraph.models.{Database, Entity, Schema}
import org.thp.scalligraph.services.{EdgeSrv, GenIntegrityCheckOps, IntegrityCheckOps, VertexSrv}
import org.thp.scalligraph.services.{EdgeSrv, GenIntegrityCheckOps, VertexSrv}
import org.thp.thehive.models._
import org.thp.thehive.services._
import play.api.Logger
Expand Down
Loading

0 comments on commit e12a821

Please sign in to comment.