from flask import Flask, request, send_file, after_this_request
import os
import subprocess
import socket

app = Flask(__name__)
app.config['UPLOAD_FOLDER'] = 'uploads'
app.config['MODIFIED_FOLDER'] = 'modified'

# Make sure the uploaded and modified folders exist
os.makedirs(app.config['UPLOAD_FOLDER'], exist_ok=True)
os.makedirs(app.config['MODIFIED_FOLDER'], exist_ok=True)

@app.route('/api/v1/pcap/comment', methods=['POST'])
def upload_file():
    if 'file' not in request.files:
        return "No file part", 400

    file = request.files['file']
    if file.filename == '':
        return "No selected file", 400

    if file:
        # Save uploaded files
        filepath = os.path.join(app.config['UPLOAD_FOLDER'], file.filename)
        file.save(filepath)

        # Use PcapNGFormatAnalys.py
        modified_filepath = modify_file_with_script(filepath)

        # Delete uploaded and modified files after response

        @after_this_request
        def remove_file(response):
            try:
                os.remove(filepath)  # 删除上传的文件
                os.remove(modified_filepath)  # 删除修改后的文件
            except Exception as e:
                app.logger.error(f'Error deleting file: {e}')
            return response

        # Provide modified file download
        return send_file(modified_filepath, as_attachment=True)

def modify_file_with_script(input_filepath):
    modified_filepath = os.path.join(app.config['MODIFIED_FOLDER'], os.path.basename(input_filepath))
    # Use PcapNGFormatAnalys.py to add Remark
    script_path = 'PcapNGFormatAnalys.py'
    subprocess.run(['python', script_path, input_filepath, modified_filepath])
    return modified_filepath

if __name__ == '__main__':
    hostname = socket.gethostname()
    local_ip = socket.gethostbyname(hostname)
    app.run(host=local_ip, port=5000, debug=True)