From 3a3eb2a44d52627ea856373774b268a7892b8587 Mon Sep 17 00:00:00 2001
From: crizzitello <98421672+crizzitello@users.noreply.github.com>
Date: Mon, 27 Jun 2022 19:27:02 -0400
Subject: [PATCH] Use QProcess to start screenshot tool (#179)

Co-authored-by: Chris Rizzitello <crizzitello@ics.com>
---
 src/helpers/screenshot.cpp | 67 ++++++++++++++++----------------------
 src/helpers/screenshot.h   | 11 ++++---
 2 files changed, 34 insertions(+), 44 deletions(-)

diff --git a/src/helpers/screenshot.cpp b/src/helpers/screenshot.cpp
index 6aea9b0..9ff47f9 100644
--- a/src/helpers/screenshot.cpp
+++ b/src/helpers/screenshot.cpp
@@ -6,12 +6,7 @@
 #include <QDir>
 #include <QFile>
 #include <QObject>
-
-#include <array>
-#include <cstdio>
-#include <iostream>
-#include <string>
-#include <utility>
+#include <QProcess>
 
 #include "appconfig.h"
 #include "helpers/file_helpers.h"
@@ -19,17 +14,6 @@
 
 Screenshot::Screenshot(QObject *parent) : QObject(parent) {}
 
-QString Screenshot::formatScreenshotCmd(QString cmdProto, const QString &filename) {
-  auto lowerCmd = cmdProto.toLower();
-  QString key = QStringLiteral("%file");
-  auto idx = lowerCmd.indexOf(key);
-  if (idx == -1) {
-    return cmdProto;
-  }
-  QString fixedFilename = QStringLiteral("'%1'").arg(filename);
-  return cmdProto.replace(idx, key.length(), fixedFilename);
-}
-
 void Screenshot::captureArea() { basicScreenshot(AppConfig::getInstance().screenshotExec); }
 
 void Screenshot::captureWindow() { basicScreenshot(AppConfig::getInstance().captureWindowExec); }
@@ -37,30 +21,35 @@ void Screenshot::captureWindow() { basicScreenshot(AppConfig::getInstance().capt
 QString Screenshot::mkName() {
   return FileHelpers::randomFilename(QStringLiteral("ashirt_screenshot_XXXXXX.%1").arg(extension()));
 }
-QString Screenshot::contentType() { return QStringLiteral("image"); }
-QString Screenshot::extension() { return QStringLiteral("png"); }
 
+void Screenshot::basicScreenshot(QString cmdProto)
+{
+    if(!QDir().mkpath(SystemHelpers::pathToEvidence()))
+        return;
+    auto newName = mkName();
+    auto tempFile = QDir::toNativeSeparators(m_fileTemplate.arg(QDir::tempPath(), newName));
+    cmdProto.replace(QStringLiteral("%file"), tempFile);
 
-void Screenshot::basicScreenshot(QString cmdProto) {
-  auto root = SystemHelpers::pathToEvidence();
-  auto hasPath = QDir().mkpath(root);
+    QString app;
+    if(cmdProto.startsWith(m_doubleQuote))
+        app = cmdProto.mid(0, cmdProto.indexOf(m_doubleQuote, 1) +1);
+    else
+        app = cmdProto.mid(0, cmdProto.indexOf(m_space));
+    cmdProto.remove(app);
+    cmdProto = cmdProto.simplified();
 
-  if (hasPath) {
-    auto tempPath = QStringLiteral("%1/%2").arg(QDir::tempPath(), mkName());
+    QProcess *ssTool = new QProcess(this);
+    ssTool->setProgram(app);
+    ssTool->setArguments(cmdProto.split(m_space));
+    ssTool->setWorkingDirectory(QDir::rootPath());
+    ssTool->start();
 
-    QString cmd = formatScreenshotCmd(std::move(cmdProto), tempPath);
-    auto lastSlash = tempPath.lastIndexOf(QStringLiteral("/")) + 1;
-    QString tempName = tempPath.right(tempPath.length() - lastSlash);
-
-    system(cmd.toStdString().c_str());
-
-    // check if file exists before doing this
-    auto finalName = root + tempName;
-    QFile src(tempPath);
-    if (src.exists()) {
-      auto moved = src.rename(QString(finalName));
-      auto trueName = moved ? finalName : tempName;
-      Q_EMIT onScreenshotCaptured(trueName);
-    }
-  }
+    connect(ssTool, &QProcess::finished, this, [this, tempFile, newName] {
+        if(!QFile::exists(tempFile))
+            return;
+        auto finalName = QDir::toNativeSeparators(SystemHelpers::pathToEvidence().append(newName));
+        auto trueName = QFile::rename(tempFile, finalName) ? finalName : newName;
+        Q_EMIT onScreenshotCaptured(trueName);
+    });
+    connect(ssTool, &QProcess::aboutToClose, ssTool, &QProcess::deleteLater);
 }
diff --git a/src/helpers/screenshot.h b/src/helpers/screenshot.h
index 3bd3c9b..ab82dba 100644
--- a/src/helpers/screenshot.h
+++ b/src/helpers/screenshot.h
@@ -4,23 +4,24 @@
 #pragma once
 
 #include <QObject>
-#include <string>
 
 class Screenshot : public QObject {
   Q_OBJECT
  public:
-  Screenshot(QObject* parent = 0);
+  Screenshot(QObject* parent = nullptr);
   void captureArea();
   void captureWindow();
 
   static QString mkName();
-  static QString extension();
-  static QString contentType();
+  static QString extension() { return QStringLiteral("png"); }
+  static QString contentType() { return QStringLiteral("image"); }
 
  signals:
   void onScreenshotCaptured(QString filepath);
 
  private:
-  QString formatScreenshotCmd(QString cmdProto, const QString& filename);
   void basicScreenshot(QString cmdProto);
+  inline static const QString m_fileTemplate = QStringLiteral("%1/%2");
+  inline static const QString m_doubleQuote = QStringLiteral("\"");
+  inline static const QString m_space = QStringLiteral(" ");
 };