/*
Hypatia: A realtime malware scanner for Android
Copyright (c) 2017-2018 Divested Computing Group

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU Affero General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU Affero General Public License for more details.

You should have received a copy of the GNU Affero General Public License
along with this program.  If not, see <https://www.gnu.org/licenses/>.
*/
package us.spotco.malwarescanner;

import android.app.Activity;
import android.app.Notification;
import android.app.NotificationChannel;
import android.app.NotificationManager;
import android.app.PendingIntent;
import android.content.Context;
import android.content.Intent;
import android.os.AsyncTask;
import android.os.Build;
import android.os.Environment;
import android.os.SystemClock;
import android.widget.TextView;

import com.google.common.hash.BloomFilter;

import java.io.File;
import java.io.FileInputStream;
import java.io.InputStream;
import java.math.BigInteger;
import java.security.MessageDigest;
import java.text.NumberFormat;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Random;
import java.util.Set;
import java.util.concurrent.ConcurrentSkipListSet;

class MalwareScanner extends AsyncTask<HashSet<File>, Object, String> {

    private final Context context;
    private TextView logOutput = null;
    private final boolean userFacing;
    private final NotificationManager notificationManager;
    private long scanStartTime = 0;
    private final HashMap<File, String> fileHashesMD5 = new HashMap<>();
    private final HashMap<File, String> fileHashesSHA1 = new HashMap<>();
    private final HashMap<File, String> fileHashesSHA256 = new HashMap<>();
    public boolean running = false;
    private int amtMatchedFiles = 0;

    public MalwareScanner(Activity activity, Context context, boolean userFacing) {
        this.context = context;
        this.userFacing = userFacing;
        if (activity != null) {
            logOutput = activity.findViewById(R.id.txtLogOutput);
        }
        notificationManager = (NotificationManager) context.getSystemService(Context.NOTIFICATION_SERVICE);
        if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.O) {
            NotificationChannel detectionChannel = new NotificationChannel("DETECTION", context.getString(R.string.lblNotificationMalwareDetectionTitle), NotificationManager.IMPORTANCE_HIGH);
            detectionChannel.setDescription(context.getString(R.string.lblNotificationMalwareDetectionDescription));
            notificationManager.createNotificationChannel(detectionChannel);
        }
    }

    private void logResult(String result, String hashsum) {
        if (result.startsWith("Potential match")) {
            String[] malwareDetect = result.split(" in ");
            if (hashsum != null) {
                //Skip if we've alerted on this file already in the past 15 seconds
                if (Utils.MATCHED_FILES_TIME.containsKey(result + hashsum)
                        && ((System.currentTimeMillis() - Utils.MATCHED_FILES_TIME.get(result + hashsum)) < (15 * 1000))) {
                    return;
                }
                Utils.MATCHED_FILES_TIME.put(result + hashsum, System.currentTimeMillis());
            }
            int notificationId = new Random().nextInt();
            Notification.Builder mBuilder =
                    new Notification.Builder(context.getApplicationContext())
                            .setSmallIcon(R.drawable.ic_notification)
                            .setContentTitle(context.getText(R.string.lblNotificationRealtimeDetection) + " " + malwareDetect[0])
                            .setContentText("File: >>>" + malwareDetect[1] + "<<<")
                            .setStyle(new Notification.BigTextStyle().bigText("File: >>>" + malwareDetect[1] + "<<<\nID: >>>" + hashsum.substring(0, 8) + "<<<"))
                            .setOngoing(true)
                            .setPriority(Notification.PRIORITY_MAX)
                            .setDefaults(Notification.DEFAULT_VIBRATE);
            if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.LOLLIPOP) {
                mBuilder.setVisibility(Notification.VISIBILITY_SECRET);
            }
            if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.O) {
                mBuilder.setChannelId("DETECTION");
            }
            //Lookup action
            Intent lookupIntent = new Intent(context, NotificationPromptActivity.class);
            lookupIntent.setFlags(Intent.FLAG_ACTIVITY_REORDER_TO_FRONT | Intent.FLAG_ACTIVITY_NEW_TASK);
            lookupIntent.setAction("us.spotco.malwarescanner.LOOKUP_HASH");
            lookupIntent.putExtra("NOTIFICATION_ID", notificationId);
            lookupIntent.putExtra("HASH", hashsum);
            PendingIntent lookupIntentPending = PendingIntent.getActivity(context, notificationId + 1, lookupIntent, PendingIntent.FLAG_IMMUTABLE);
            mBuilder.addAction(android.R.drawable.ic_dialog_map, context.getText(R.string.lookupVT), lookupIntentPending);
            //Delete action
            if (malwareDetect[1].startsWith("~/")) {
                Intent deleteIntent = new Intent(context, NotificationPromptActivity.class);
                deleteIntent.setFlags(Intent.FLAG_ACTIVITY_REORDER_TO_FRONT | Intent.FLAG_ACTIVITY_NEW_TASK);
                deleteIntent.setAction("us.spotco.malwarescanner.DELETE_FILE");
                deleteIntent.putExtra("NOTIFICATION_ID", notificationId);
                deleteIntent.putExtra("FILE_PATH", malwareDetect[1].replaceFirst("~", Environment.getExternalStorageDirectory().toString()));
                PendingIntent deleteIntentPending = PendingIntent.getActivity(context, notificationId + 2, deleteIntent, PendingIntent.FLAG_IMMUTABLE);
                mBuilder.addAction(android.R.drawable.ic_delete, context.getText(R.string.deleteFile), deleteIntentPending);
            } else if (malwareDetect[1].startsWith("/data/app")) {
                Intent uninstallIntent = new Intent(context, NotificationPromptActivity.class);
                uninstallIntent.setFlags(Intent.FLAG_ACTIVITY_REORDER_TO_FRONT | Intent.FLAG_ACTIVITY_NEW_TASK);
                uninstallIntent.setAction("us.spotco.malwarescanner.UNINSTALL_APP");
                uninstallIntent.putExtra("NOTIFICATION_ID", notificationId);
                uninstallIntent.putExtra("FILE_PATH", malwareDetect[1]);
                PendingIntent uninstallIntentPending = PendingIntent.getActivity(context, notificationId + 3, uninstallIntent, PendingIntent.FLAG_IMMUTABLE);
                mBuilder.addAction(android.R.drawable.ic_delete, context.getText(R.string.uninstallApp), uninstallIntentPending);
            }
            //Ignore action
            Intent ignoreIntent = new Intent(context, NotificationPromptActivity.class);
            ignoreIntent.setFlags(Intent.FLAG_ACTIVITY_REORDER_TO_FRONT | Intent.FLAG_ACTIVITY_NEW_TASK);
            ignoreIntent.setAction("us.spotco.malwarescanner.IGNORE_DETECTION");
            ignoreIntent.putExtra("NOTIFICATION_ID", notificationId);
            PendingIntent ignoreIntentPending = PendingIntent.getActivity(context, notificationId + 4, ignoreIntent, PendingIntent.FLAG_IMMUTABLE);
            mBuilder.addAction(android.R.drawable.ic_menu_view, context.getText(R.string.ignoreDetection), ignoreIntentPending);

            //Show it!
            notificationManager.notify(notificationId, mBuilder.build());
        } else if (userFacing) {
            logOutput.append(result);
            if (!(result.length() <= 3)) {
                logOutput.append("\n");
            }
        }
    }

    @Override
    protected final void onPreExecute() {
        scanStartTime = SystemClock.elapsedRealtime();
        logResult(context.getString(R.string.main_starting_scan), null);
    }

    @Override
    protected final String doInBackground(HashSet<File>[] filesToScan) {
        running = true;
        int startCount = amtMatchedFiles;
        ConcurrentSkipListSet<File> filesToScanReal = new ConcurrentSkipListSet<>(); //TODO: Reduce this?
        for (Set<File> fileArray : filesToScan) {
            for (File file : fileArray) {
                filesToScanReal.addAll(Utils.getFilesRecursive(file)); //TODO: Inline this, hash files as they are found
            }
        }

        //Pre
        fileHashesMD5.clear();
        fileHashesSHA1.clear();
        fileHashesSHA256.clear();

        publishProgress("\t" + context.getString(R.string.main_files_pending_scan, NumberFormat.getInstance().format(filesToScanReal.size())) + "\n", true);

        Database.loadDatabase(context, false, Database.signatureDatabases);
        int delayCount = 0;
        if (Database.areDatabasesAvailable()) {
            while (!Database.isDatabaseLoaded() && delayCount <= 90) {
                try {
                    Thread.sleep(1000);
                    delayCount++;
                    if ((delayCount % 15) == 0) {
                        publishProgress("\t" + context.getString(R.string.main_database_loading), true);
                    }
                } catch (InterruptedException e) {
                    e.printStackTrace();
                }
            }
        }
        if (Database.isDatabaseLoaded()) {
            publishProgress("\t" + context.getString(R.string.main_database_loaded, NumberFormat.getInstance().format(Database.signaturesCount)) + "\n", true);

            //Perform a self-test
            if (Database.selfTest()) {
                publishProgress("\t" + context.getString(R.string.self_test_result_success), true);
            } else {
                publishProgress("\t" + context.getString(R.string.self_test_result_failure), true);
            }

            //Get file hashes
            publishProgress("\t" + context.getString(R.string.main_hashing_files), true);
            publishProgress("\t", true);
            int fileScannedCount = 0;
            int percentIncrement = (filesToScanReal.size() / 20);
            if (percentIncrement < 1) { //Prevent divide by zero
                percentIncrement = 1;
            }
            String spinnerCur = " ~ ";
            long totalBytesHashed = 0;
            long hashStartTime = SystemClock.elapsedRealtime();
            for (File file : filesToScanReal) {
                if (this.isCancelled()) { //Allow quicker cancels
                    //publishProgress("\t" + context.getString(R.string.main_cancelled_scan), true);
                    running = false;
                    return null;
                }
                totalBytesHashed += file.length();
                getFileHashes(file);
                filesToScanReal.remove(file);
                fileScannedCount++;
                if ((fileScannedCount % percentIncrement) == 0) {
                    publishProgress(spinnerCur, true);
                    if (spinnerCur.equals(" = ")) {
                        spinnerCur = " ~ ";
                    } else {
                        spinnerCur = " = ";
                    }
                }
                //Log.d("Hypatia", "Scanning " + file);
            }
            filesToScanReal.clear();
            publishProgress("\n\t" + context.getString(R.string.main_hashing_done) + "\n", true);

            //Check the hashes
            checkSignature("MD5", fileHashesMD5, Database.signaturesMD5);
            checkSignature("SHA-1", fileHashesSHA1, Database.signaturesSHA1);
            checkSignature("SHA-256", fileHashesSHA256, Database.signaturesSHA256);
            if (Database.signaturesMD5Extended != null) {
                checkSignature("MD5 Extended", fileHashesMD5, Database.signaturesMD5Extended);
            }

            //Scan result
            if (startCount == amtMatchedFiles) {
                publishProgress("\n\t" + context.getString(R.string.detections_none) + "\n", true);
            } else {
                publishProgress("\n\t" + context.getString(R.string.detections_found) + "\n", true);
            }

            //Post
            fileHashesMD5.clear();
            fileHashesSHA1.clear();
            fileHashesSHA256.clear();
            Utils.FILES_SCANNED.getAndAdd(fileScannedCount);
            if (userFacing || Utils.FILES_SCANNED.get() % 40 == 0) {
                System.gc(); //GC can be expensive, don't run it too often.
            }
            if (userFacing) {
                long secondsSpent = ((SystemClock.elapsedRealtime() - scanStartTime) / 1000L);
                long secondsSpentHashing = ((SystemClock.elapsedRealtime() - hashStartTime) / 1000L);
                long MBS = 0;
                if (secondsSpentHashing > 0) {
                    MBS = totalBytesHashed / 1000 / 1000 / secondsSpentHashing;
                }
                publishProgress(context.getString(R.string.main_scanning_done, String.valueOf(secondsSpent), String.valueOf(MBS)) + "\n\n\n\n", true);
            }
        } else {
            publishProgress("\t" + context.getString(R.string.main_no_database_available), true);
        }
        running = false;
        return null;
    }

    @Override
    protected final void onProgressUpdate(Object... objects) {
        String hash = null;
        if (objects.length == 3) {
            hash = (String) objects[2];
        }
        logResult((String) objects[0], hash);
    }

    private void checkSignature(String hashType, HashMap<File, String> signaturesToCheck, BloomFilter<String> signatureDatabase) {
        if (Database.isDatabaseLoaded() && signatureDatabase != null) {
            for (Map.Entry<File, String> file : signaturesToCheck.entrySet()) {
                if (!file.getValue().equals("d41d8cd98f00b204e9800998ecf8427e")
                        && !file.getValue().equals("da39a3ee5e6b4b0d3255bfef95601890afd80709")
                        && !file.getValue().equals("e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855")) {
                    if (signatureDatabase.mightContain(file.getValue())) {
                        amtMatchedFiles++;
                        //Log.d("Hypatia", "Match: " + fileHashesSHA256.get(file.getKey()));
                        publishProgress("Potential match in " + file.getKey().toString().replaceAll(Environment.getExternalStorageDirectory().toString(), "~"), false, fileHashesSHA256.get(file.getKey()));
                    } else {
                        //Log.d("Hypatia", "No match for " + file.getValue());
                    }
                }
            }
            publishProgress("\t" + context.getString(R.string.main_hash_scan_done, hashType) + "\n", true);
        } else {
            publishProgress("\t" + context.getString(R.string.main_no_database_available) + "\n", true);
        }
    }

    private void getFileHashes(File file) {
        try {
            InputStream fis = new FileInputStream(file);

            byte[] buffer = new byte[4096];
            int numRead;

            MessageDigest digestMD5 = MessageDigest.getInstance("MD5");
            MessageDigest digestSHA1 = MessageDigest.getInstance("SHA-1");
            MessageDigest digestSHA256 = MessageDigest.getInstance("SHA-256");

            do {
                numRead = fis.read(buffer);
                if (numRead > 0) {
                    digestMD5.update(buffer, 0, numRead);
                    digestSHA1.update(buffer, 0, numRead);
                    digestSHA256.update(buffer, 0, numRead);
                }
            } while (numRead != -1);

            fis.close();

            fileHashesMD5.put(file, String.format("%032x", new BigInteger(1, digestMD5.digest())).toLowerCase());
            fileHashesSHA1.put(file, String.format("%040x", new BigInteger(1, digestSHA1.digest())).toLowerCase());
            fileHashesSHA256.put(file, String.format("%064x", new BigInteger(1, digestSHA256.digest())).toLowerCase());
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

}