flask_wtf の FileFiled から受け取ったファイルの中身をハッシュ化したい

Question

flask_wtfのFileFieldから受け取ったファイルの中身をハッシュ化してデータベースに登録しようとしたのですが，どのファイルを選択しても登録されるハッシュ値が同じになってしまいます．おそらくFileStorage object に原因があるのだと思うのですが，わかる方は教えていただけないでしょうか．FileFieldではファイルの中身を読み込めないのであれば，どのようにすればよいか教えてください．よろしくお願いいたします．

必要なimportは全て行われており，エラーは出ていないのでそこには問題はないと思います．

データを受け取るFileFieldはこれになります．

class MalwareRegistrationForm(FlaskForm):
    malware_name = StringField('Malware Name', validators=[DataRequired()])
    malware_file = FileField('Choose Malware file', validators=[FileRequired()])
    submit = SubmitField('Register')

    def validate_malware_file(self, malware_file):
        target = generate_file_hash(self.malware_file.data)

        if Malware.query.filter_by(malware_hash=target).first():
            raise ValidationError('Malware hash has been registerd already!')

受け取ったファイルのハッシュ値を計算する関数はこれです．

def generate_file_hash(file):

    file_content = file.read()
    file_content = hashlib.sha256(file_content)
    return file_content.hexdigest()

ファイルを登録するデータベースのクラスがこれです．

class Malware(db.Model, UserMixin):

    __tablename__ = 'malware'

    id = db.Column(db.Integer, primary_key = True)
    malware_name = db.Column(db.String(64), index=True)
    malware_hash = db.Column(db.String(128))

    # ファイルを受け取ってハッシュ化してのち登録
    def __init__(self, malware_name, malware_file):
        self.malware_name = malware_name
        self.malware_hash = generate_file_hash(malware_file)

そして，ファイルの選択及びアップロード画面のview.py部分がこれです

@check_blueprint.route('/addmalware', methods=['GET', 'POST'])
@login_required
def malware_register():
    form = MalwareRegistrationForm()

    if form.validate_on_submit():
        malware = Malware(malware_name=form.malware_name.data,
                          malware_file=form.malware_file.data)

        db.session.add(malware)
        db.session.commit()
        flash('Successfully registerd.')
        return redirect(url_for('index'))

    return render_template('register_malware.html', form=form)

なにか少しでも気付いたことがありましたら教えていただけると幸いです．

Answer 1

結局FileFieldの使用は諦め，以下のコードにすることで解決しました．

ALLOWED_EXTENSIONS = {'txt', 'pdf', 'png', 'jpg', 'jpeg', 'exe'}

def allowed_file(filename):
    return '.' in filename and \
           filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS

@check_blueprint.route('/check', methods=['GET', 'POST'])
def upload_file():

    flag = ''
    name = ''

    if request.method == 'POST':
        # check if the post request has the file part
        if 'file' not in request.files:
            flash('No file part')
        file = request.files['file']
        # if user does not select file, browser also
        # submit an empty part without filename
        filename = file.filename

        if filename == '':
            flash('No selected file')
        if file and allowed_file(filename):
            filename = secure_filename(filename)
            filepath = os.path.join(app.config['UPLOAD_FOLDER'], filename)
            file.save(filepath)

            hash_val = generate_file_hash(filepath)
            malware = Malware.query.filter_by(malware_hash=hash_val).first()

            os.remove(filepath)

            if malware is not None:
                flag = 'True'
                name = malware.malware_name
            else:
                flag = 'False'

            # result = History(is_malware=str(flag), file_hash=hash_val)
            result = History(is_malware=flag, hash_val=hash_val)

            db.session.add(result)
            db.session.commit()

        elif allowed_file(filename) == False:
            flash("Your file's extension is not accepted.")
        else:
            pass



    return render_template('select_file.html', flag=flag, name=name)

def generate_file_hash(filepath):

    with open(filepath) as file:
        binary = file.read()

    hash_val = hashlib.sha256(binary.encode('utf-8')).hexdigest()

    return hash_val

ご協力ありがとうございました．

Answer 2

以下のような単純な例でハッシュ値を計算できました。
コード全体が分からないですが、正常にアップロードされたファイルデータを取得できていないのではないでしょうか。

from flask import Flask, render_template
from flask_wtf import FlaskForm
from flask_wtf.file import FileField, FileRequired
from wtforms import SubmitField
import hashlib

class FileForm(FlaskForm):
    file = FileField(validators=[FileRequired()])
    submit = SubmitField('Upload')

app = Flask(__name__)
app.config['SECRET_KEY'] = 'hoge'

@app.route("/", methods=['GET', 'POST'])
def upload():
    form = FileForm()
    if form.validate_on_submit():
        f = form.file.data
        data = f.read()
        hash = hashlib.sha256(data)
        print(data, hash.hexdigest())
        # b'abc' ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad

    return render_template('index.html', form=form)

if __name__ == "__main__":
    app.run(debug=True)

index.html

<!doctype html><html lang="ja">
<head><meta charset="UTF-8"><title>アップロード</title></head>
<body>
<form method="POST" enctype="multipart/form-data">
     {{ form.hidden_tag() }}
     {{ form.file }}
     {{ form.submit }}
</form>
</body>
</html>

実行結果

 * Serving Flask app "sv" (lazy loading)
 * Environment: production
   WARNING: This is a development server. Do not use it in a production deployment.
   Use a production WSGI server instead.
 * Debug mode: on
 * Restarting with windowsapi reloader
 * Debugger is active!
 * Debugger PIN: 256-310-102
 * Running on http://127.0.0.1:5000/ (Press CTRL+C to quit)
127.0.0.1 - - [04/Aug/2020 13:04:34] "GET / HTTP/1.1" 200 -
127.0.0.1 - - [04/Aug/2020 13:04:34] "GET /favicon.ico HTTP/1.1" 404 -
b'abc' ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad
127.0.0.1 - - [04/Aug/2020 13:04:55] "POST / HTTP/1.1" 200 -
b'def' cb8379ac2098aa165029e3938a51da0bcecfc008fd6795f401178647f96c5b34
127.0.0.1 - - [04/Aug/2020 13:05:16] "POST / HTTP/1.1" 200 -

Comments

[tomiharu1017]

回答ありがとうございます．私も初めにこのコードを試したのですが，この場合だと全てのファイルに対して同じハッシュ値を返してしまうのです．読み込んだ内容をprint()してみたのですが，中身は空でした．

[8524ba23]

回答コードを実際に動作させてみましたか？正常に動作します。
コメントにあるような現象を示す実行結果を提示ください。

[tomiharu1017]

今試したところcan110さんのコードが正常に動作しました！コードのロジックは私のものと同じはずなのですが...
しかしFileFieldでもできることが分かったので，もう一度確認してみます．
本当にありがとうございました．

[tomiharu1017]

ところで，このようにform.file.dataのようにしてハッシュ化するコードはセキュリティの観点からして安全と言えるのでしょうか？もし何か知見がありましたら教えていただけると幸いです．厚かましくて申し訳ありません．

[8524ba23]

何をもって安全とみなすかによるかと思いますが、任意のバイナリからハッシュ値を求める処理自体には問題はないように思います。

[tomiharu1017]

了解しました．丁寧な対応ありがとうございました．