From 163e73651585cc97c29e1f27561c06a251fd8a91 Mon Sep 17 00:00:00 2001
From: To-om <toom@thehive-project.org>
Date: Thu, 4 Apr 2019 15:22:37 +0200
Subject: [PATCH] #181 Add organization in OAuth2 user mapper

---
 .../cortex/services/mappers/GroupUserMapper.scala   | 11 ++++++++++-
 .../cortex/services/mappers/SimpleUserMapper.scala  | 13 +++++++++++--
 2 files changed, 21 insertions(+), 3 deletions(-)

diff --git a/app/org/thp/cortex/services/mappers/GroupUserMapper.scala b/app/org/thp/cortex/services/mappers/GroupUserMapper.scala
index 2ead6580c..3c940d4d4 100644
--- a/app/org/thp/cortex/services/mappers/GroupUserMapper.scala
+++ b/app/org/thp/cortex/services/mappers/GroupUserMapper.scala
@@ -16,7 +16,9 @@ class GroupUserMapper(
     nameAttrName: String,
     rolesAttrName: Option[String],
     groupAttrName: String,
+    organizationAttrName: Option[String],
     defaultRoles: Seq[String],
+    defaultOrganization: Option[String],
     groupsUrl: String,
     mappings: Map[String, Seq[String]],
     ws: WSClient,
@@ -31,7 +33,9 @@ class GroupUserMapper(
     configuration.getOptional[String]("auth.sso.attributes.name").getOrElse("username"),
     configuration.getOptional[String]("auth.sso.attributes.roles"),
     configuration.getOptional[String]("auth.sso.attributes.groups").getOrElse(""),
+    configuration.getOptional[String]("auth.sso.attributes.organization"),
     configuration.getOptional[Seq[String]]("auth.sso.defaultRoles").getOrElse(Seq()),
+    configuration.getOptional[String]("auth.sso.defaultOrganization"),
     configuration.getOptional[String]("auth.sso.groups.url").getOrElse(""),
     configuration.getOptional[Map[String, Seq[String]]]("auth.sso.groups.mappings").getOrElse(Map()),
     ws,
@@ -50,10 +54,15 @@ class GroupUserMapper(
       val fields = for {
         login ← (jsValue \ loginAttrName).validate[String]
         name ← (jsValue \ nameAttrName).validate[String]
+        organization ← organizationAttrName
+          .flatMap(o ⇒ (jsValue \ o).asOpt[String])
+          .orElse(defaultOrganization)
+          .fold[JsResult[String]](JsError())(o ⇒ JsSuccess(o))
       } yield Fields(Json.obj(
         "login" → login,
         "name" → name,
-        "roles" → roles))
+        "roles" → roles,
+        "organization" → organization))
       fields match {
         case JsSuccess(f, _) ⇒ Future.successful(f)
         case JsError(errors) ⇒ Future.failed(AuthenticationError(s"User info fails: ${errors.map(_._1).mkString}"))
diff --git a/app/org/thp/cortex/services/mappers/SimpleUserMapper.scala b/app/org/thp/cortex/services/mappers/SimpleUserMapper.scala
index 03ffc4588..b7e6919fe 100644
--- a/app/org/thp/cortex/services/mappers/SimpleUserMapper.scala
+++ b/app/org/thp/cortex/services/mappers/SimpleUserMapper.scala
@@ -3,7 +3,7 @@ package org.thp.cortex.services.mappers
 import scala.concurrent.{ ExecutionContext, Future }
 
 import play.api.Configuration
-import play.api.libs.json.{ JsError, JsSuccess, JsValue, Json }
+import play.api.libs.json._
 
 import javax.inject.Inject
 
@@ -14,14 +14,18 @@ class SimpleUserMapper(
     loginAttrName: String,
     nameAttrName: String,
     rolesAttrName: Option[String],
+    organizationAttrName: Option[String],
     defaultRoles: Seq[String],
+    defaultOrganization: Option[String],
     implicit val ec: ExecutionContext) extends UserMapper {
 
   @Inject() def this(configuration: Configuration, ec: ExecutionContext) = this(
     configuration.getOptional[String]("auth.sso.attributes.login").getOrElse("name"),
     configuration.getOptional[String]("auth.sso.attributes.name").getOrElse("username"),
     configuration.getOptional[String]("auth.sso.attributes.roles"),
+    configuration.getOptional[String]("auth.sso.attributes.organization"),
     configuration.getOptional[Seq[String]]("auth.sso.defaultRoles").getOrElse(Seq()),
+    configuration.getOptional[String]("auth.sso.defaultOrganization"),
     ec)
 
   override val name: String = "simple"
@@ -31,10 +35,15 @@ class SimpleUserMapper(
       login ← (jsValue \ loginAttrName).validate[String]
       name ← (jsValue \ nameAttrName).validate[String]
       roles = rolesAttrName.fold(defaultRoles)(r ⇒ (jsValue \ r).asOpt[Seq[String]].getOrElse(defaultRoles))
+      organization ← organizationAttrName
+        .flatMap(o ⇒ (jsValue \ o).asOpt[String])
+        .orElse(defaultOrganization)
+        .fold[JsResult[String]](JsError())(o ⇒ JsSuccess(o))
     } yield Fields(Json.obj(
       "login" → login,
       "name" → name,
-      "roles" → roles))
+      "roles" → roles,
+      "organization" → organization))
     fields match {
       case JsSuccess(f, _) ⇒ Future.successful(f)
       case JsError(errors) ⇒ Future.failed(AuthenticationError(s"User info fails: ${errors.map(_._1).mkString}"))