Keine Beschreibung

case_assets_db.py 11KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421
  1. # IRIS Source Code
  2. # Copyright (C) 2021 - Airbus CyberSecurity (SAS)
  3. # ir@cyberactionlab.net
  4. #
  5. # This program is free software; you can redistribute it and/or
  6. # modify it under the terms of the GNU Lesser General Public
  7. # License as published by the Free Software Foundation; either
  8. # version 3 of the License, or (at your option) any later version.
  9. #
  10. # This program is distributed in the hope that it will be useful,
  11. # but WITHOUT ANY WARRANTY; without even the implied warranty of
  12. # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
  13. # Lesser General Public License for more details.
  14. #
  15. # You should have received a copy of the GNU Lesser General Public License
  16. # along with this program; if not, write to the Free Software Foundation,
  17. # Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
  18. import datetime
  19. from flask_login import current_user
  20. from sqlalchemy import and_
  21. from sqlalchemy import func
  22. from flask_sqlalchemy.pagination import Pagination
  23. from app import db, app
  24. from app.datamgmt.states import update_assets_state
  25. from app.datamgmt.conversions import convert_sort_direction
  26. from app.models.models import AnalysisStatus
  27. from app.models.models import CaseStatus
  28. from app.models.models import AssetComments
  29. from app.models.models import AssetsType
  30. from app.models.models import CaseAssets
  31. from app.models.models import CaseEventsAssets
  32. from app.models.cases import Cases
  33. from app.models.models import Comments
  34. from app.models.models import CompromiseStatus
  35. from app.models.models import Ioc
  36. from app.models.models import IocAssetLink
  37. from app.models.models import IocType
  38. from app.models.authorization import User
  39. from app.models.pagination_parameters import PaginationParameters
  40. log = app.logger
  41. def create_asset(asset, caseid, user_id):
  42. asset.date_added = datetime.datetime.utcnow()
  43. asset.date_update = datetime.datetime.utcnow()
  44. asset.case_id = caseid
  45. asset.user_id = user_id
  46. db.session.add(asset)
  47. update_assets_state(caseid=caseid, userid=user_id)
  48. db.session.commit()
  49. return asset
  50. def case_assets_db_exists(asset: CaseAssets):
  51. assets = CaseAssets.query.filter(
  52. func.lower(CaseAssets.asset_name) == func.lower(asset.asset_name),
  53. CaseAssets.asset_type_id == asset.asset_type_id,
  54. CaseAssets.asset_id != asset.asset_id,
  55. CaseAssets.case_id == asset.case_id
  56. )
  57. return assets.first() is not None
  58. def get_assets(case_identifier):
  59. assets = CaseAssets.query.with_entities(
  60. CaseAssets.asset_id,
  61. CaseAssets.asset_uuid,
  62. CaseAssets.asset_name,
  63. AssetsType.asset_name.label('asset_type'),
  64. AssetsType.asset_icon_compromised,
  65. AssetsType.asset_icon_not_compromised,
  66. CaseAssets.asset_description,
  67. CaseAssets.asset_domain,
  68. CaseAssets.asset_compromise_status_id,
  69. CaseAssets.asset_ip,
  70. CaseAssets.asset_type_id,
  71. AnalysisStatus.name.label('analysis_status'),
  72. CaseAssets.analysis_status_id,
  73. CaseAssets.asset_tags
  74. ).filter(
  75. CaseAssets.case_id == case_identifier,
  76. ).join(
  77. CaseAssets.asset_type
  78. ).join(
  79. CaseAssets.analysis_status
  80. ).all()
  81. return assets
  82. def filter_assets(case_identifier, pagination_parameters: PaginationParameters) -> Pagination:
  83. query = CaseAssets.query.filter(
  84. CaseAssets.case_id == case_identifier
  85. )
  86. order_by = pagination_parameters.get_order_by()
  87. if order_by is not None:
  88. order_func = convert_sort_direction(pagination_parameters.get_direction())
  89. if hasattr(CaseAssets, order_by):
  90. query = query.order_by(order_func(getattr(CaseAssets, order_by)))
  91. assets = query.paginate(page=pagination_parameters.get_page(), per_page=pagination_parameters.get_per_page(), error_out=False)
  92. return assets
  93. def get_raw_assets(caseid):
  94. assets = CaseAssets.query.filter(
  95. CaseAssets.case_id == caseid
  96. ).all()
  97. return assets
  98. def get_assets_name(caseid):
  99. assets_names = CaseAssets.query.with_entities(
  100. CaseAssets.asset_name
  101. ).filter(
  102. CaseAssets.case_id == caseid
  103. ).all()
  104. return assets_names
  105. def get_asset(asset_id) -> CaseAssets:
  106. asset = CaseAssets.query.filter(
  107. CaseAssets.asset_id == asset_id,
  108. ).first()
  109. return asset
  110. def update_asset(asset_name, asset_description, asset_ip, asset_info, asset_domain,
  111. asset_compromise_status_id, asset_type, asset_id, caseid, analysis_status, asset_tags):
  112. asset = get_asset(asset_id)
  113. asset.asset_name = asset_name
  114. asset.asset_description = asset_description
  115. asset.asset_ip = asset_ip
  116. asset.asset_info = asset_info
  117. asset.asset_domain = asset_domain
  118. asset.asset_compromise_status_id = asset_compromise_status_id
  119. asset.asset_type_id = asset_type
  120. asset.analysis_status_id = analysis_status
  121. asset.asset_tags = asset_tags
  122. update_assets_state(caseid=caseid)
  123. db.session.commit()
  124. def delete_asset(asset: CaseAssets):
  125. delete_ioc_asset_link(asset.asset_id)
  126. # Delete the relevant records from the CaseEventsAssets table
  127. CaseEventsAssets.query.filter(
  128. asset.asset_id == CaseEventsAssets.asset_id
  129. ).delete()
  130. # Delete the relevant records from the AssetComments table
  131. com_ids = AssetComments.query.with_entities(
  132. AssetComments.comment_id
  133. ).filter(
  134. AssetComments.comment_asset_id == asset.asset_id
  135. ).all()
  136. com_ids = [c.comment_id for c in com_ids]
  137. AssetComments.query.filter(AssetComments.comment_id.in_(com_ids)).delete()
  138. Comments.query.filter(
  139. Comments.comment_id.in_(com_ids)
  140. ).delete()
  141. db.session.delete(asset)
  142. update_assets_state(asset.case_id)
  143. db.session.commit()
  144. def get_assets_types():
  145. assets_types = [(c.asset_id, c.asset_name) for c
  146. in AssetsType.query.with_entities(AssetsType.asset_name,
  147. AssetsType.asset_id).order_by(AssetsType.asset_name)
  148. ]
  149. return assets_types
  150. def get_unspecified_analysis_status_id():
  151. """
  152. Get the id of the 'Unspecified' analysis status
  153. """
  154. analysis_status = AnalysisStatus.query.filter(
  155. AnalysisStatus.name == 'Unspecified'
  156. ).first()
  157. return analysis_status.id if analysis_status else None
  158. def get_analysis_status_list():
  159. analysis_status = [(c.id, c.name) for c in AnalysisStatus.query.with_entities(
  160. AnalysisStatus.id,
  161. AnalysisStatus.name
  162. )]
  163. return analysis_status
  164. def get_compromise_status_list():
  165. return [(e.value, e.name.replace('_', ' ').capitalize()) for e in CompromiseStatus]
  166. def get_compromise_status_dict():
  167. return [{'value': e.value, 'name': e.name.replace('_', ' ').capitalize()} for e in CompromiseStatus]
  168. def get_case_outcome_status_dict():
  169. return [{'value': e.value, 'name': e.name.replace('_', ' ').capitalize()} for e in CaseStatus]
  170. def get_asset_type_id(asset_type_name):
  171. assets_type_id = AssetsType.query.with_entities(
  172. AssetsType.asset_id
  173. ).filter(
  174. func.lower(AssetsType.asset_name) == asset_type_name
  175. ).first()
  176. return assets_type_id
  177. def get_assets_ioc_links(caseid):
  178. ioc_links_req = IocAssetLink.query.with_entities(
  179. Ioc.ioc_id,
  180. Ioc.ioc_value,
  181. IocAssetLink.asset_id
  182. ).filter(
  183. Ioc.ioc_id == IocAssetLink.ioc_id,
  184. Ioc.case_id == caseid
  185. ).all()
  186. return ioc_links_req
  187. def get_similar_assets(asset_name, asset_type_id, caseid, customer_id, cases_limitation):
  188. linked_assets = CaseAssets.query.with_entities(
  189. Cases.name.label('case_name'),
  190. Cases.open_date.label('case_open_date'),
  191. CaseAssets.asset_description,
  192. CaseAssets.asset_compromise_status_id,
  193. CaseAssets.asset_id,
  194. CaseAssets.case_id
  195. ).filter(
  196. Cases.client_id == customer_id,
  197. CaseAssets.case_id != caseid
  198. ).filter(
  199. CaseAssets.asset_name == asset_name,
  200. CaseAssets.asset_type_id == asset_type_id,
  201. Cases.case_id.in_(cases_limitation)
  202. ).join(CaseAssets.case).all()
  203. return (lasset._asdict() for lasset in linked_assets)
  204. def delete_ioc_asset_link(asset_id):
  205. IocAssetLink.query.filter(
  206. IocAssetLink.asset_id == asset_id
  207. ).delete()
  208. def get_linked_iocs_from_asset(asset_id):
  209. iocs = IocAssetLink.query.with_entities(
  210. Ioc.ioc_id,
  211. Ioc.ioc_value
  212. ).filter(
  213. IocAssetLink.asset_id == asset_id,
  214. Ioc.ioc_id == IocAssetLink.ioc_id
  215. ).all()
  216. return iocs
  217. def set_ioc_links(ioc_list, asset_id):
  218. if ioc_list is None:
  219. return False, "Empty IOC list"
  220. # Reset IOC list
  221. delete_ioc_asset_link(asset_id)
  222. for ioc in ioc_list:
  223. ial = IocAssetLink()
  224. ial.asset_id = asset_id
  225. ial.ioc_id = ioc
  226. db.session.add(ial)
  227. try:
  228. db.session.commit()
  229. except Exception as e:
  230. db.session.rollback()
  231. log.exception(e)
  232. return True, e.__str__()
  233. return False, ""
  234. def get_linked_iocs_id_from_asset(asset_id):
  235. iocs = IocAssetLink.query.with_entities(
  236. IocAssetLink.ioc_id
  237. ).filter(
  238. IocAssetLink.asset_id == asset_id
  239. ).all()
  240. return iocs
  241. def get_linked_iocs_finfo_from_asset(asset_id):
  242. iocs = IocAssetLink.query.with_entities(
  243. Ioc.ioc_id,
  244. Ioc.ioc_value,
  245. Ioc.ioc_tags,
  246. Ioc.ioc_type_id,
  247. IocType.type_name,
  248. Ioc.ioc_description,
  249. Ioc.ioc_tlp_id
  250. ).filter(and_(
  251. IocAssetLink.asset_id == asset_id,
  252. IocAssetLink.ioc_id == Ioc.ioc_id
  253. )).join(Ioc.ioc_type).all()
  254. return iocs
  255. def get_case_asset_comments(asset_id):
  256. return Comments.query.filter(
  257. AssetComments.comment_asset_id == asset_id
  258. ).with_entities(
  259. Comments
  260. ).join(AssetComments,
  261. Comments.comment_id == AssetComments.comment_id
  262. ).order_by(
  263. Comments.comment_date.asc()
  264. ).all()
  265. def add_comment_to_asset(asset_id, comment_id):
  266. ec = AssetComments()
  267. ec.comment_asset_id = asset_id
  268. ec.comment_id = comment_id
  269. db.session.add(ec)
  270. db.session.commit()
  271. def get_case_assets_comments_count(asset_id):
  272. return AssetComments.query.filter(
  273. AssetComments.comment_asset_id.in_(asset_id)
  274. ).with_entities(
  275. AssetComments.comment_asset_id,
  276. AssetComments.comment_id
  277. ).group_by(
  278. AssetComments.comment_asset_id,
  279. AssetComments.comment_id
  280. ).all()
  281. def get_case_asset_comment(asset_id, comment_id):
  282. return AssetComments.query.filter(
  283. AssetComments.comment_asset_id == asset_id,
  284. AssetComments.comment_id == comment_id
  285. ).with_entities(
  286. Comments.comment_id,
  287. Comments.comment_text,
  288. Comments.comment_date,
  289. Comments.comment_update_date,
  290. Comments.comment_uuid,
  291. User.name,
  292. User.user
  293. ).join(
  294. AssetComments.comment
  295. ).join(
  296. Comments.user
  297. ).first()
  298. def delete_asset_comment(asset_id, comment_id, case_id):
  299. comment = Comments.query.filter(
  300. Comments.comment_id == comment_id,
  301. Comments.comment_user_id == current_user.id
  302. ).first()
  303. if not comment:
  304. return False, "You are not allowed to delete this comment"
  305. AssetComments.query.filter(
  306. AssetComments.comment_asset_id == asset_id,
  307. AssetComments.comment_id == comment_id
  308. ).delete()
  309. db.session.delete(comment)
  310. db.session.commit()
  311. return True, "Comment deleted"
  312. def get_asset_by_name(asset_name, caseid):
  313. asset = CaseAssets.query.filter(
  314. CaseAssets.asset_name == asset_name,
  315. CaseAssets.case_id == caseid
  316. ).first()
  317. return asset