11import os
22import shutil
33from model import *
4- from utils import user_assigned_assessment
4+ from utils import *
55from werkzeug .utils import secure_filename
6- from flask import Blueprint , request , send_from_directory , jsonify
6+ from flask import Blueprint , request , send_from_directory , jsonify , make_response
77from flask_security import auth_required , roles_accepted , current_user
88
99blueprint_testcase_utils = Blueprint ('blueprint_testcase_utils' , __name__ )
1212@auth_required ()
1313@roles_accepted ('Admin' , 'Red' )
1414@user_assigned_assessment
15- def testcasevisibility (id ):
16- newcase = TestCase .objects (id = id ).first ()
17- newcase .visible = not newcase .visible
18- newcase .save ()
19-
20- return jsonify (newcase .to_json ()), 200
15+ def toggle_visibility_flag (id ):
16+ testcase = get_testcase_by_id (id )
17+ if not testcase :
18+ return "Test case not found" , 404
19+
20+ testcase .visible = not testcase .visible
21+ testcase .save ()
22+ return jsonify (testcase .to_json ()), 200
23+
2124
2225@blueprint_testcase_utils .route ('/testcase/<id>/clone' , methods = ['POST' ])
2326@auth_required ()
2427@roles_accepted ('Admin' , 'Red' )
2528@user_assigned_assessment
26- def testcaseclone (id ):
27- orig = TestCase .objects (id = id ).first ()
28- newcase = TestCase ()
29- copy = ["name" , "assessmentid" , "objective" , "requirements" , "actions" , "rednotes" , "mitreid" , "tactic" , "tools" , "tags" , "expectedprevention" , "expectedalertcreation" , "expectedincidentcreation" , "expectedseverity" , "priorityurgency" ]
30- for field in copy :
31- newcase [field ] = orig [field ]
32- newcase .name = orig ["name" ] + " (Copy)"
33- newcase .save ()
29+ def clone_testcase (id ):
30+ orig = get_testcase_by_id (id )
31+ if not orig :
32+ return "Test case not found" , 404
3433
34+ fields_to_copy = [
35+ "name" , "assessmentid" , "objective" , "requirements" , "actions" , "rednotes" ,
36+ "mitreid" , "tactic" , "tools" , "tags" , "expectedprevention" ,
37+ "expectedalertcreation" , "expectedincidentcreation" , "expectedseverity" , "priorityurgency"
38+ ]
39+ newcase = TestCase (** {field : orig [field ] for field in fields_to_copy })
40+ newcase .name = f"{ orig ['name' ]}
41+ newcase .save ()
3542 return jsonify (newcase .to_json ()), 200
3643
44+
3745@blueprint_testcase_utils .route ('/testcase/<id>/toggle-delete' , methods = ['POST' ])
3846@auth_required ()
3947@roles_accepted ('Admin' , 'Red' )
4048@user_assigned_assessment
41- def testcasedeleted (id ):
42- newcase = TestCase .objects (id = id ).first ()
43- newcase .deleted = not newcase .deleted
44- newcase .save ()
45-
49+ def toggle_delete_flag (id ):
50+ testcase = get_testcase_by_id (id )
51+ if not testcase :
52+ return "Test case not found" , 404
53+
54+ testcase .deleted = not testcase .deleted
55+ testcase .save ()
4656 return "" , 200
4757
58+
4859@blueprint_testcase_utils .route ('/testcase/<id>/delete' , methods = ['POST' ])
4960@auth_required ()
5061@roles_accepted ('Admin' )
5162@user_assigned_assessment
52- def testcasedelete (id ):
53- testcase = TestCase .objects (id = id ).first ()
63+ def delete_testcase (id ):
64+ testcase = get_testcase_by_id (id )
65+ if not testcase :
66+ return "Test case not found" , 404
67+
5468 assessment = Assessment .objects (id = testcase .assessmentid ).first ()
5569 if os .path .exists (f"files/{ str (assessment .id )} { str (testcase .id )} ):
5670 shutil .rmtree (f"files/{ str (assessment .id )} { str (testcase .id )} )
5771 testcase .delete ()
5872
59- return "" , 200
73+ return "Test case deleted" , 200
74+
6075
6176@blueprint_testcase_utils .route ('/testcase/<id>/evidence/<colour>/<file>' , methods = ['DELETE' ])
6277@auth_required ()
6378@roles_accepted ('Admin' , 'Red' , 'Blue' )
6479@user_assigned_assessment
65- def deletefile (id , colour , file ):
66- if colour not in "red" , "blue" ]:
67- return 401
80+ def delete_file (id , colour , file ):
81+ VALID_COLOURS = {'red' , 'blue' }
82+
83+ if colour not in VALID_COLOURS :
84+ return "Invalid colour" , 400
6885 if colour == "red" and current_user .has_role ("Blue" ):
69- return 403
86+ return "Invalid colour" , 400
7087
71- testcase = TestCase .objects (id = id ).first ()
72- # Sanity check to prevent death if the image has already been removed
73- path = f"files/{ testcase .assessmentid } { testcase .id } { secure_filename (file )}
74- if os .path .isfile (path ):
75- os .remove (path )
76-
77- files = []
78- for f in testcase ["redfiles" if colour == "red" else "bluefiles" ]:
79- if f .name != file :
80- files .append (f )
81-
82- if colour == "red" :
83- testcase .update (set__redfiles = files )
84- else :
85- testcase .update (set__bluefiles = files )
86-
87- return '' , 204
88-
89- @blueprint_testcase_utils .route ('/testcase/<id>/evidence/<file>' , methods = ['GET' ])
88+ testcase = get_testcase_by_id (id )
89+ if not testcase :
90+ return "Test case not found" , 404
91+
92+ filepath = f"files/{ testcase .assessmentid } { testcase .id } { secure_filename (file )}
93+ if os .path .isfile (filepath ):
94+ os .remove (filepath )
95+
96+ filelist = testcase .redfiles if colour == "red" else testcase .bluefiles
97+ updated_files = [f for f in filelist if f .name != file ]
98+
99+ update_field = 'redfiles' if colour == "red" else 'bluefiles'
100+ testcase .update (** {f"set__{ update_field } : updated_files })
101+
102+ return "" , 204
103+
104+
105+ @blueprint_testcase_utils .route ('/testcase/<id>/evidence/<file>' , methods = ['GET' ])
90106@auth_required ()
91107@user_assigned_assessment
92- def fetchFile (id , file ):
93- testcase = TestCase .objects (id = id ).first ()
94-
95- return send_from_directory (
96- 'files' ,
97- f"{ testcase .assessmentid } { str (testcase .id )} { secure_filename (file )} ,
98- as_attachment = True if "download" in request .args else False
99- )
108+ def fetch_file (id , file ):
109+ ALLOWED_INLINE_EXTENSIONS = {'.png' , '.jpg' , '.jpeg' }
110+
111+ testcase = get_testcase_by_id (id )
112+ if not testcase :
113+ return "Test case not found" , 404
114+
115+ filename = secure_filename (file )
116+ folder_path = os .path .join ('files' , str (testcase .assessmentid ), str (testcase .id ))
117+ file_path = os .path .join (folder_path , filename )
118+
119+ if not os .path .isfile (file_path ):
120+ return "File not found" , 404
121+
122+ _ , ext = os .path .splitext (filename )
123+ ext = ext .lower ()
124+ as_attachment = not (ext in ALLOWED_INLINE_EXTENSIONS and "download" not in request .args )
125+
126+ response = make_response (send_from_directory (folder_path , filename , as_attachment = as_attachment ))
127+ response .headers ["X-Content-Type-Options" ] = "nosniff"
128+ return response
0 commit comments