Use QProcess to start screenshot tool (#179)

Co-authored-by: Chris Rizzitello <crizzitello@ics.com>
main
crizzitello 2022-06-27 19:27:02 -04:00 committed by GitHub
parent d5ed38d747
commit 3a3eb2a44d
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 34 additions and 44 deletions

View File

@ -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);
}

View File

@ -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(" ");
};