generated from SAP/repository-template
-
Notifications
You must be signed in to change notification settings - Fork 7
Release v0.6.2 #138
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Merged
Release v0.6.2 #138
Changes from all commits
Commits
Show all changes
13 commits
Select commit
Hold shift + click to select a range
89f716b
re organized routes + service file to clarify and ease run_all whethe…
ea63453
Bump actions/setup-node from 5 to 6
dependabot[bot] df995b9
Bump the js-dependencies group in /frontend with 17 updates
dependabot[bot] 0a3850c
Use python3.12 instead of 3.11
marcorosa 5c653a6
Do not remove docs folders
marcorosa 0a9c0a9
Merge pull request #137 from SAP/dependabot/npm_and_yarn/frontend/dev…
marcorosa e7efc62
Merge pull request #136 from SAP/dependabot/github_actions/develop/ac…
marcorosa 0857f74
Improve validation and error handling in routes and services
c080837
forgot a abort(403) in verify_api_key to match the behavior
73cdcc7
fixed line too long for linter
ba82360
Merge pull request #125 from SAP/re-organize-backend
marcorosa b577f01
Bump versions
marcorosa 56002d8
[Changelog CI] Add Changelog for Version v0.6.2
github-actions[bot] File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,245 @@ | ||
| import json | ||
| import os | ||
|
|
||
| from flask import request, jsonify, send_file, abort | ||
| from sqlalchemy import select | ||
|
|
||
| from app.db.models import Attack, ModelAttackScore, TargetModel, db | ||
| from app.utils import send_intro, verify_api_key | ||
| from attack_result import SuiteResult | ||
| from services import run_all_attacks | ||
| from status import status | ||
|
|
||
|
|
||
| def register_routes(app, sock, agent=None, callbacks=None): | ||
| # ---------------------- | ||
| # Health endpoints | ||
| # ---------------------- | ||
| @app.route("/health") | ||
| def check_health(): | ||
| """ | ||
| Health route is used in the CI to test that the installation was | ||
| successful. | ||
| """ | ||
| return jsonify({'status': 'ok'}) | ||
|
|
||
| # ---------------------- | ||
| # Attacks endpoints | ||
| # ---------------------- | ||
| @app.route('/run_all', methods=['POST']) | ||
| def execute_all_attacks(): | ||
| """ | ||
| Run all attacks. Used for automation. | ||
| Expected JSON body: | ||
| { | ||
| "target": "string" | ||
| } | ||
| """ | ||
| verify_api_key() | ||
| data = request.get_json() | ||
| target_model = data.get('target') if data else None | ||
| if not target_model: | ||
| return jsonify({'error': 'target parameter is required'}), 400 | ||
| # Call the service to run all attacks | ||
| result = run_all_attacks( | ||
| target=target_model | ||
| ) | ||
| return jsonify(result), 200 if result.get('success') else 500 | ||
|
|
||
| @app.route('/api/attacks', methods=['GET']) | ||
| def get_attacks(): | ||
| """ | ||
| Endpoint to retrieve all attacks with their weights. | ||
| Returns a JSON object with attack names and their weights. | ||
| """ | ||
| try: | ||
| attacks = db.session.query(Attack).all() | ||
| attack_list = [ | ||
| {'name': attack.name, 'weight': attack.weight} | ||
| for attack in attacks | ||
| ] | ||
| return jsonify(attack_list), 200 | ||
| except Exception as e: | ||
| return jsonify({'error': str(e)}), 500 | ||
|
|
||
| @app.route('/api/attacks', methods=['PUT']) | ||
| def update_attack_weights(): | ||
| """ | ||
| Update weights for multiple attacks. | ||
| Expects a JSON object like: {"artPrompt": 2, "codeAttack": 1, ...} | ||
| """ | ||
| verify_api_key() | ||
| try: | ||
| weights = request.get_json() | ||
| if not isinstance(weights, dict): | ||
| return jsonify({'error': 'Invalid payload format'}), 400 | ||
|
|
||
| for name, weight in weights.items(): | ||
| attack = db.session.query(Attack).filter_by(name=name).first() | ||
| if attack: | ||
| attack.weight = float(weight) | ||
| else: | ||
| return jsonify({'error': f'Attack not found: {name}'}), 404 | ||
|
|
||
| db.session.commit() | ||
| return jsonify({'message': 'Weights updated successfully'}), 200 | ||
|
|
||
| except Exception as e: | ||
| db.session.rollback() | ||
| return jsonify({'error': str(e)}), 500 | ||
|
|
||
| # ---------------------- | ||
| # Reports endpoints | ||
| # ---------------------- | ||
| @app.route('/download_report') | ||
| def download_report(): | ||
| """ | ||
| This route allows to download attack suite reports by specifying | ||
| their name. | ||
| """ | ||
| name = request.args.get('name') | ||
| format = request.args.get('format', 'md') | ||
|
|
||
| # Ensure that a name is provided | ||
| if not name: | ||
| abort(400) | ||
| # Ensure that only allowed chars are in the filename | ||
| # (e.g. no path traversal) | ||
| if not all([c in SuiteResult.FILENAME_ALLOWED_CHARS for c in name]): | ||
| abort(400) | ||
marcorosa marked this conversation as resolved.
Show resolved
Hide resolved
|
||
|
|
||
| results = SuiteResult.load_from_name(name) | ||
|
|
||
| generated_name = name + '_generated' | ||
| path = os.path.join(SuiteResult.DEFAULT_OUTPUT_PATH, generated_name) | ||
| result_path = results.to_file(path, format) | ||
| return send_file( | ||
| result_path, | ||
| mimetype=SuiteResult.get_mime_type(format) | ||
| ) | ||
|
|
||
| @app.route('/api/heatmap', methods=['GET']) | ||
| def get_heatmap(): | ||
| """ | ||
| Endpoint to retrieve heatmap data showing model score | ||
| against various attacks. | ||
|
|
||
| Queries the database for total attacks and successes per target model | ||
| and attack combination. | ||
| Calculates attack success rate and returns structured data for | ||
| visualization. | ||
|
|
||
| Returns: | ||
| JSON response with: | ||
| - models: List of target models and their attack success rate | ||
| per attack. | ||
| - attacks: List of attack names and their associated weights. | ||
|
|
||
| HTTP Status Codes: | ||
| 200: Data successfully retrieved. | ||
| 500: Internal server error during query execution. | ||
| """ | ||
| try: | ||
| query = ( | ||
| select( | ||
| ModelAttackScore.total_number_of_attack, | ||
| ModelAttackScore.total_success, | ||
| TargetModel.name.label('attack_model_name'), | ||
| Attack.name.label('attack_name'), | ||
| Attack.weight.label('attack_weight') | ||
| ) | ||
| .join(TargetModel, ModelAttackScore.target_model_id == TargetModel.id) # noqa: E501 | ||
| .join(Attack, ModelAttackScore.attack_id == Attack.id) | ||
| ) | ||
|
|
||
| scores = db.session.execute(query).all() | ||
| all_models = {} | ||
| all_attacks = {} | ||
|
|
||
| for score in scores: | ||
| model_name = score.attack_model_name | ||
| attack_name = score.attack_name | ||
|
|
||
| if attack_name not in all_attacks: | ||
| all_attacks[attack_name] = score.attack_weight | ||
|
|
||
| if model_name not in all_models: | ||
| all_models[model_name] = { | ||
| 'name': model_name, | ||
| 'scores': {}, | ||
| } | ||
|
|
||
| # Compute attack success rate for this model/attack | ||
| success_ratio = ( | ||
| round((score.total_success / score.total_number_of_attack) * 100) # noqa: E501 | ||
| if score.total_number_of_attack else 0 | ||
| ) | ||
|
|
||
| all_models[model_name]['scores'][attack_name] = success_ratio | ||
|
|
||
| return jsonify({ | ||
| 'models': list(all_models.values()), | ||
| 'attacks': [ | ||
| {'name': name, 'weight': weight} | ||
| for name, weight in sorted(all_attacks.items()) | ||
| ] | ||
| }) | ||
| except Exception as e: | ||
| return jsonify({'error': str(e)}), 500 | ||
|
|
||
| # ---------------------- | ||
| # WebSocket endpoints | ||
| # ---------------------- | ||
| @sock.route('/agent') | ||
| def query_agent(sock): | ||
| """ | ||
| Websocket route for the frontend to send prompts to the agent and | ||
| receive responses as well as status updates. | ||
|
|
||
| Messages received are in this JSON format: | ||
| { | ||
| "type":"message", | ||
| "data":"Start the vulnerability scan", | ||
| "key":"secretapikey" | ||
| } | ||
| """ | ||
| # Verify API key from headers before establishing session | ||
| verify_api_key() | ||
| if not agent: | ||
| sock.send(json.dumps({ | ||
| 'type': 'message', | ||
| 'data': 'Agent is disabled on this deployment.' | ||
| })) | ||
| return | ||
| status.sock = sock | ||
| # Intro is sent after connecting successfully | ||
| send_intro(sock) | ||
| while True: | ||
| try: | ||
| data_raw = sock.receive() | ||
| data = json.loads(data_raw) | ||
| assert 'data' in data | ||
| query = data['data'] | ||
| status.clear_report() | ||
| response = agent.invoke( | ||
| {'input': query}, | ||
| config=callbacks or {} | ||
| ) | ||
| ai_response = response['output'] | ||
| formatted_output = { | ||
| 'type': 'message', | ||
| 'data': ( | ||
| f'{ai_response}' | ||
| ) | ||
| } | ||
| sock.send(json.dumps(formatted_output)) | ||
| except json.JSONDecodeError: | ||
| sock.send(json.dumps({ | ||
| 'type': 'error', | ||
| 'data': 'Invalid JSON format' | ||
| })) | ||
| except Exception as e: | ||
| sock.send(json.dumps({ | ||
| 'type': 'error', | ||
| 'data': f'Error: {str(e)}' | ||
| })) | ||
marcorosa marked this conversation as resolved.
Show resolved
Hide resolved
|
||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,34 @@ | ||
| import json | ||
| import os | ||
|
|
||
| from flask import request, abort | ||
|
|
||
|
|
||
| def send_intro(sock): | ||
| """ | ||
| Sends the intro via the websocket connection. | ||
|
|
||
| The intro is meant as a short tutorial on how to use the agent. | ||
| Also it includes meaningful suggestions for prompts that should | ||
| result in predictable behavior for the agent, e.g. | ||
| "Start the vulnerability scan". | ||
| """ | ||
| intro_file = 'data/intro.txt' | ||
| try: | ||
| with open(intro_file, 'r') as f: | ||
| intro = f.read() | ||
| except FileNotFoundError: | ||
| intro = "Welcome! (intro file missing)" | ||
| sock.send(json.dumps({'type': 'message', 'data': intro})) | ||
|
|
||
|
|
||
| def verify_api_key(): | ||
| """ | ||
| Verifies the API key from the request headers against the env variable. | ||
| If no API key is configured, access is allowed. | ||
| If API key is configured but missing/invalid, request is rejected. | ||
| """ | ||
| if os.getenv('API_KEY'): | ||
| provided_key = request.headers.get('X-API-Key') | ||
| if provided_key != os.getenv('API_KEY'): | ||
| abort(403) | ||
marcorosa marked this conversation as resolved.
Show resolved
Hide resolved
|
||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Uh oh!
There was an error while loading. Please reload this page.