From c766b3e048a3c69388d74f13170f699125f09f61 Mon Sep 17 00:00:00 2001
From: lupoDharkael <izhe@hotmail.es>
Date: Tue, 8 May 2018 21:23:09 +0200
Subject: [PATCH] Add capture individual screens option

---
 README.md                         |  8 ++++
 dbus/org.dharkael.Flameshot.xml   | 20 +++++++-
 docs/bash-completion/flameshot    |  9 +++-
 src/cli/commandlineparser.cpp     |  6 ++-
 src/core/capturerequest.cpp       |  9 +++-
 src/core/capturerequest.h         |  7 ++-
 src/core/controller.cpp           | 41 +++++++++++++----
 src/core/controller.h             |  3 ++
 src/core/flameshotdbusadapter.cpp | 17 ++++++-
 src/core/flameshotdbusadapter.h   |  3 +-
 src/main.cpp                      | 76 ++++++++++++++++++++++++++++---
 src/utils/screengrabber.cpp       | 28 ++++++++++++
 src/utils/screengrabber.h         |  1 +
 13 files changed, 200 insertions(+), 28 deletions(-)

diff --git a/README.md b/README.md
index a6707630..4b42e820 100644
--- a/README.md
+++ b/README.md
@@ -58,6 +58,14 @@ Example commands:
 
    `flameshot full -c -p ~/myStuff/captures`
 
+- capture the screen containing the mouse and print the image (bytes) in PNG format:
+
+   `flameshot screen -r`
+
+- capture the screen number 1 and copy it to the clipboard:
+
+   `flameshot screen -n 1 -c`
+
 In case of doubt choose the first or the second command as shortcut in your favorite desktop environment.
 
 A systray icon will be in your system's panel while Flameshot is running.
diff --git a/dbus/org.dharkael.Flameshot.xml b/dbus/org.dharkael.Flameshot.xml
index 57b899a3..87b68b0a 100644
--- a/dbus/org.dharkael.Flameshot.xml
+++ b/dbus/org.dharkael.Flameshot.xml
@@ -18,7 +18,7 @@
     </method>
 
     <!--
-        fullScreenRaw:
+        fullScreen:
         @path: the path where the screenshot will be saved. When the argument is empty the program will ask for a path graphically.
         @toClipboard: Whether to copy the screenshot to clipboard or not.
         @delay: delay time in milliseconds, both return the @id defined in the call of this method.
@@ -33,6 +33,24 @@
       <arg name="id" type="i" direction="in"/>
     </method>
 
+    <!--
+        captureScreen:
+        @number: number of the screen to be captured.
+        @path: the path where the screenshot will be saved. When the argument is empty the program will ask for a path graphically.
+        @toClipboard: Whether to copy the screenshot to clipboard or not.
+        @delay: delay time in milliseconds, both return the @id defined in the call of this method.
+        @id: identificator of the call.
+
+        Takes a screenshot of the whole screen and sends a captureTaken signal with the raw image or a captureFailed signal.
+    -->
+    <method name="captureScreen">
+      <arg name="number" type="i" direction="in"/>
+      <arg name="path" type="s" direction="in"/>
+      <arg name="toClipboard" type="b" direction="in"/>
+      <arg name="delay" type="i" direction="in"/>
+      <arg name="id" type="i" direction="in"/>
+    </method>
+
     <!--
         openConfig:
 
diff --git a/docs/bash-completion/flameshot b/docs/bash-completion/flameshot
index 86fac301..2fe7fbe2 100644
--- a/docs/bash-completion/flameshot
+++ b/docs/bash-completion/flameshot
@@ -7,12 +7,17 @@ _flameshot() {
 
 	prev="${COMP_WORDS[COMP_CWORD-1]}"
 	cur="${COMP_WORDS[COMP_CWORD]}"
-	cmd="gui full config"
+	cmd="gui full config screen"
+	screen_opts="--number --path --delay --raw -p -d -r -n"
 	gui_opts="--path --delay --raw -p -d -r"
 	full_opts="--path --delay --clipboard --raw -p -d -c -r"
 	config_opts="--contrastcolor --filename --maincolor --showhelp --trayicon -k -f -m -s -t"
 
 	case "${prev}" in
+		screen)
+			COMPREPLY=( $(compgen -W "$screen_opts --help -h" -- ${cur}) )
+			return 0
+			;;
 		gui)
 			COMPREPLY=( $(compgen -W "$gui_opts --help -h" -- ${cur}) )
 			return 0
@@ -33,7 +38,7 @@ _flameshot() {
 			COMPREPLY=( $(compgen -W "true false" -- ${cur}) )
 			return 0
 			;;
-		-d|--delay|-h|--help|-c|--clipboard|--version|-v)
+		-d|--delay|-h|--help|-c|--clipboard|--version|-v|--number|-n)
 			return 0
 			;;
 		*)
diff --git a/src/cli/commandlineparser.cpp b/src/cli/commandlineparser.cpp
index 0ff5a628..ebcc3e31 100644
--- a/src/cli/commandlineparser.cpp
+++ b/src/cli/commandlineparser.cpp
@@ -62,10 +62,12 @@ QString optionsToString(const QList<CommandOption> &options,
     QString result;
     if(!dashedOptionList.isEmpty()) {
         result += "Options:\n";
+        QString linePadding = QString(" ").repeated(size + 4).prepend("\n");
         for (int i = 0; i < options.length(); ++i) {
             result += QStringLiteral("  %1  %2\n")
                     .arg(dashedOptionList.at(i).leftJustified(size, ' '))
-                    .arg(options.at(i).description());
+                    .arg(options.at(i).description()
+                         .replace("\n", linePadding));
         }
         if (!arguments.isEmpty()) {
             result += "\n";
@@ -299,7 +301,7 @@ bool CommandLineParser::isSet(const CommandOption &option) const {
 }
 
 QString CommandLineParser::value(const CommandOption &option) const {
-    QString value;
+    QString value = option.value();
     for (const CommandOption &fOption: m_foundOptions) {
         if (option == fOption) {
             value = fOption.value();
diff --git a/src/core/capturerequest.cpp b/src/core/capturerequest.cpp
index d08927b2..99bbc2e6 100644
--- a/src/core/capturerequest.cpp
+++ b/src/core/capturerequest.cpp
@@ -22,9 +22,10 @@
 
 CaptureRequest::CaptureRequest(CaptureRequest::CaptureMode mode,
                                const uint delay, const QString &path,
+                               const QVariant &data,
                                CaptureRequest::ExportTask tasks) :
     m_mode(mode), m_delay(delay), m_path(path), m_tasks(tasks),
-    m_forcedID(false), m_id(0)
+    m_data(data), m_forcedID(false), m_id(0)
 {
 
 }
@@ -42,7 +43,7 @@ uint CaptureRequest::id() const {
     uint id = 0;
     QVector<uint>v;
     v << qHash(m_mode) << qHash(m_delay * QDateTime::currentMSecsSinceEpoch())
-      << qHash(m_path) << qHash(m_tasks);
+      << qHash(m_path) << qHash(m_tasks) << m_data.toInt();
     for(uint i : v) {
         id ^= i + 0x9e3779b9 + (id << 6) + (id >> 2);
     }
@@ -61,6 +62,10 @@ QString CaptureRequest::path() const {
     return m_path;
 }
 
+QVariant CaptureRequest::data() const {
+    return m_data;
+}
+
 void CaptureRequest::addTask(CaptureRequest::ExportTask task) {
     m_tasks |= task;
 }
diff --git a/src/core/capturerequest.h b/src/core/capturerequest.h
index 39abdcfb..c0d6814b 100644
--- a/src/core/capturerequest.h
+++ b/src/core/capturerequest.h
@@ -19,14 +19,14 @@
 
 #include <QString>
 #include <QPixmap>
+#include <QVariant>
 
 class CaptureRequest {
 public:
     enum CaptureMode {
         FULLSCREEN_MODE,
         GRAPHICAL_MODE,
-        //GRAPHICAL_WINDOW,
-        //SCREEN,
+        SCREEN_MODE,
     };
 
     enum ExportTask {
@@ -38,6 +38,7 @@ public:
     CaptureRequest(CaptureMode mode,
                    const uint delay = 0,
                    const QString &path = "",
+                   const QVariant &data = QVariant(),
                    ExportTask tasks = NO_TASK);
 
     void setStaticID(uint id);
@@ -45,6 +46,7 @@ public:
     uint id() const;
     uint delay() const;
     QString path() const;
+    QVariant data() const;
     CaptureMode captureMode() const;
 
     void addTask(ExportTask task);
@@ -55,6 +57,7 @@ private:
     uint m_delay;
     QString m_path;
     ExportTask m_tasks;
+    QVariant m_data;
 
     bool m_forcedID;
     uint m_id;
diff --git a/src/core/controller.cpp b/src/core/controller.cpp
index 651db8b1..12c83d17 100644
--- a/src/core/controller.cpp
+++ b/src/core/controller.cpp
@@ -28,6 +28,7 @@
 #include <QSystemTrayIcon>
 #include <QAction>
 #include <QMenu>
+#include <QDesktopWidget>
 
 #ifdef Q_OS_WIN
 #include "src/core/globalshortcutfilter.h"
@@ -57,11 +58,6 @@ Controller::Controller() : m_captureWindow(nullptr) {
 
     QString StyleSheet = CaptureButton::globalStyleSheet();
     qApp->setStyleSheet(StyleSheet);
-
-    connect(this, &Controller::captureTaken,
-            this, &Controller::handleCaptureTaken);
-    connect(this, &Controller::captureFailed,
-            this, &Controller::handleCaptureFailed);
 }
 
 Controller *Controller::getInstance() {
@@ -69,6 +65,13 @@ Controller *Controller::getInstance() {
     return &c;
 }
 
+void Controller::enableExports() {
+    connect(this, &Controller::captureTaken,
+            this, &Controller::handleCaptureTaken);
+    connect(this, &Controller::captureFailed,
+            this, &Controller::handleCaptureFailed);
+}
+
 void Controller::requestCapture(const CaptureRequest &request) {
     uint id = request.id();
     m_requestMap.insert(id, request);
@@ -79,7 +82,13 @@ void Controller::requestCapture(const CaptureRequest &request) {
             this->startFullscreenCapture(id);
         });
         break;
-    case CaptureRequest::GRAPHICAL_MODE: {
+    case CaptureRequest::SCREEN_MODE: {
+        int &&number = request.data().toInt();
+        doLater(request.delay(), this, [this, id, number](){
+            this->startScreenGrab(id, number);
+        });
+        break;
+    } case CaptureRequest::GRAPHICAL_MODE: {
         QString &&path = request.path();
         doLater(request.delay(), this, [this, id, path](){
             this->startVisualCapture(id, path);
@@ -121,6 +130,22 @@ void Controller::startVisualCapture(const uint id, const QString &forcedSavePath
     }
 }
 
+void Controller::startScreenGrab(const uint id, const int screenNumber) {
+    bool ok = true;
+    int n = screenNumber;
+
+    if (n < 0) {
+        QPoint globalCursorPos = QCursor::pos();
+        n = qApp->desktop()->screenNumber(globalCursorPos);
+    }
+    QPixmap p(ScreenGrabber().grabScreen(n, ok));
+    if (ok) {
+        emit captureTaken(id, p);
+    } else {
+        emit captureFailed(id);
+    }
+}
+
 // creation of the configuration window
 void Controller::openConfigWindow() {
     if (!m_configWindow) {
@@ -226,10 +251,6 @@ void Controller::handleCaptureFailed(uint id) {
 }
 
 void Controller::doLater(int msec, QObject *receiver, lambda func)  {
-    if (msec == 0) {
-        func();
-        return;
-    }
     QTimer *timer = new QTimer(receiver);
     QObject::connect(timer, &QTimer::timeout, receiver,
                      [timer, func](){ func(); timer->deleteLater(); });
diff --git a/src/core/controller.h b/src/core/controller.h
index b58d9b1d..17ee7f0a 100644
--- a/src/core/controller.h
+++ b/src/core/controller.h
@@ -41,6 +41,8 @@ public:
     Controller(const Controller&) = delete;
     void operator =(const Controller&) = delete;
 
+    void enableExports();
+
 signals:
     void captureTaken(uint id, QPixmap p);
     void captureFailed(uint id);
@@ -62,6 +64,7 @@ private slots:
     void startFullscreenCapture(const uint id = 0);
     void startVisualCapture(const uint id = 0,
                              const QString &forcedSavePath = QString());
+    void startScreenGrab(const uint id = 0, const int screenNumber = -1);
 
     void handleCaptureTaken(uint id, QPixmap p);
     void handleCaptureFailed(uint id);
diff --git a/src/core/flameshotdbusadapter.cpp b/src/core/flameshotdbusadapter.cpp
index e9133bbd..92b85d8d 100644
--- a/src/core/flameshotdbusadapter.cpp
+++ b/src/core/flameshotdbusadapter.cpp
@@ -22,7 +22,6 @@
 #include "src/utils/screenshotsaver.h"
 #include "src/utils/systemnotification.h"
 #include <QBuffer>
-
 FlameshotDBusAdapter::FlameshotDBusAdapter(QObject *parent)
     : QDBusAbstractAdaptor(parent)
 {
@@ -60,6 +59,20 @@ void FlameshotDBusAdapter::fullScreen(
     Controller::getInstance()->requestCapture(req);
 }
 
+void FlameshotDBusAdapter::captureScreen(int number, QString path,
+                                         bool toClipboard, int delay, uint id)
+{
+    CaptureRequest req(CaptureRequest::SCREEN_MODE, delay, path, number);
+    if (toClipboard) {
+        req.addTask(CaptureRequest::CLIPBOARD_SAVE_TASK);
+    }
+    if (!path.isEmpty()) {
+        req.addTask(CaptureRequest::FILESYSTEM_SAVE_TASK);
+    }
+    req.setStaticID(id);
+    Controller::getInstance()->requestCapture(req);
+}
+
 void FlameshotDBusAdapter::openConfig() {
     Controller::getInstance()->openConfigWindow();
 }
@@ -80,7 +93,7 @@ void FlameshotDBusAdapter::autostartEnabled(bool enabled) {
     controller->updateConfigComponents();
 }
 
-void FlameshotDBusAdapter::handleCaptureTaken(uint id, QPixmap p) {
+void FlameshotDBusAdapter::handleCaptureTaken(uint id, const QPixmap &p) {
     QByteArray byteArray;
     QBuffer buffer(&byteArray);
     p.save(&buffer, "PNG");
diff --git a/src/core/flameshotdbusadapter.h b/src/core/flameshotdbusadapter.h
index 5d51bb77..6fbd2228 100644
--- a/src/core/flameshotdbusadapter.h
+++ b/src/core/flameshotdbusadapter.h
@@ -35,10 +35,11 @@ signals:
 public slots:
     Q_NOREPLY void graphicCapture(QString path, int delay, uint id);
     Q_NOREPLY void fullScreen(QString path, bool toClipboard, int delay, uint id);
+    Q_NOREPLY void captureScreen(int number, QString path, bool toClipboard, int delay, uint id);
     Q_NOREPLY void openConfig();
     Q_NOREPLY void trayIconEnabled(bool enabled);
     Q_NOREPLY void autostartEnabled(bool enabled);
 
 private slots:
-    void handleCaptureTaken(uint id, QPixmap p);
+    void handleCaptureTaken(uint id, const QPixmap &p);
 };
diff --git a/src/main.cpp b/src/main.cpp
index 27bc9fb0..cf07d070 100644
--- a/src/main.cpp
+++ b/src/main.cpp
@@ -63,8 +63,8 @@ int main(int argc, char *argv[]) {
         app.setApplicationName("flameshot");
         app.setOrganizationName("Dharkael");
 
-#ifdef Q_OS_LINUX
         auto c = Controller::getInstance();
+#ifdef Q_OS_LINUX
         new FlameshotDBusAdapter(c);
         QDBusConnection dbus = QDBusConnection::sessionBus();
         if (!dbus.isConnected()) {
@@ -73,10 +73,10 @@ int main(int argc, char *argv[]) {
         }
         dbus.registerObject("/", c);
         dbus.registerService("org.dharkael.Flameshot");
-#else
-        // Create inicial static instance
-        Controller::getInstance();
 #endif
+        // Exporting captures must be connected after the dbus interface
+        // or the dbus signal gets blocked until we end the exports.
+        c->enableExports();
         return app.exec();
     }
 
@@ -97,6 +97,7 @@ int main(int argc, char *argv[]) {
     CommandArgument fullArgument("full", "Capture the entire desktop.");
     CommandArgument guiArgument("gui", "Start a manual capture in GUI mode.");
     CommandArgument configArgument("config", "Configure flameshot.");
+    CommandArgument screenArgument("screen", "Capture a single screen.");
 
     // Options
     CommandOption pathOption(
@@ -137,6 +138,10 @@ int main(int argc, char *argv[]) {
     CommandOption rawImageOption(
                 {"r", "raw"},
                 "Print raw PNG capture");
+    CommandOption screenNumberOption(
+                {"n", "number"},
+                "Define the screen to capture,\ndefault: screen containing the cursor",
+                "Screen number", "-1");
 
     // Add checkers
     auto colorChecker = [&parser](const QString &colorCode) -> bool {
@@ -152,7 +157,8 @@ int main(int argc, char *argv[]) {
                        "You may need to escape the '#' sign as in '\\#FFF'";
 
     const QString delayErr = "Invalid delay, it must be higher than 0";
-    auto delayChecker = [&parser](const QString &delayValue) -> bool {
+    const QString numberErr = "Invalid screen number, it must be non negative";
+    auto numericChecker = [&parser](const QString &delayValue) -> bool {
         int value = delayValue.toInt();
         return value >= 0;
     };
@@ -173,19 +179,24 @@ int main(int argc, char *argv[]) {
 
     contrastColorOption.addChecker(colorChecker, colorErr);
     mainColorOption.addChecker(colorChecker, colorErr);
-    delayOption.addChecker(delayChecker, delayErr);
+    delayOption.addChecker(numericChecker, delayErr);
     pathOption.addChecker(pathChecker, pathErr);
     trayOption.addChecker(booleanChecker, booleanErr);
     autostartOption.addChecker(booleanChecker, booleanErr);
     showHelpOption.addChecker(booleanChecker, booleanErr);
+    screenNumberOption.addChecker(numericChecker, numberErr);
 
     // Relationships
     parser.AddArgument(guiArgument);
+    parser.AddArgument(screenArgument);
     parser.AddArgument(fullArgument);
     parser.AddArgument(configArgument);
     auto helpOption = parser.addHelpOption();
     auto versionOption = parser.addVersionOption();
     parser.AddOptions({ pathOption, delayOption, rawImageOption }, guiArgument);
+    parser.AddOptions({ screenNumberOption, clipboardOption,pathOption,
+                        delayOption, rawImageOption },
+                      screenArgument);
     parser.AddOptions({ pathOption, clipboardOption, delayOption, rawImageOption },
                       fullArgument);
     parser.AddOptions({ autostartOption, filenameOption, trayOption,
@@ -248,6 +259,9 @@ int main(int argc, char *argv[]) {
         if (toClipboard) {
             req.addTask(CaptureRequest::CLIPBOARD_SAVE_TASK);
         }
+        if (!pathValue.isEmpty()) {
+            req.addTask(CaptureRequest::FILESYSTEM_SAVE_TASK);
+        }
         uint id = req.id();
         DBusUtils dbusUtils;
 
@@ -271,6 +285,56 @@ int main(int argc, char *argv[]) {
             app.exec();
         }
     }
+    else if (parser.isSet(screenArgument)) { // SCREEN
+        QString numberStr = parser.value(screenNumberOption);
+        int number = numberStr.startsWith("-") ? -1 : numberStr.toInt();
+        QString pathValue = parser.value(pathOption);
+        int delay = parser.value(delayOption).toInt();
+        bool toClipboard = parser.isSet(clipboardOption);
+        bool isRaw = parser.isSet(rawImageOption);
+        // Not a valid command
+        if (!isRaw && !toClipboard && pathValue.isEmpty()) {
+            QTextStream out(stdout);
+            out << "Invalid format, set where to save the content with one of "
+                << "the following flags:\n "
+                << pathOption.dashedNames().join(", ") << "\n "
+                << rawImageOption.dashedNames().join(", ") << "\n "
+                << clipboardOption.dashedNames().join(", ") << "\n\n";
+            parser.parse(QStringList() << argv[0] << "screen" << "-h");
+            goto finish;
+        }
+
+        CaptureRequest req(CaptureRequest::SCREEN_MODE,
+                           delay, pathValue, number);
+        if (toClipboard) {
+            req.addTask(CaptureRequest::CLIPBOARD_SAVE_TASK);
+        }
+        if (!pathValue.isEmpty()) {
+            req.addTask(CaptureRequest::FILESYSTEM_SAVE_TASK);
+        }
+        uint id = req.id();
+        DBusUtils dbusUtils;
+
+        // Send message
+        QDBusMessage m = QDBusMessage::createMethodCall("org.dharkael.Flameshot",
+                                               "/", "", "captureScreen");
+        m << number << pathValue << toClipboard << delay << id;
+        QDBusConnection sessionBus = QDBusConnection::sessionBus();
+        dbusUtils.checkDBusConnection(sessionBus);
+        sessionBus.call(m);
+
+        if (isRaw) {
+            dbusUtils.connectPrintCapture(sessionBus, id);
+            // timeout just in case
+            QTimer t;
+            t.setInterval(delay + 2000);
+            QObject::connect(&t, &QTimer::timeout, qApp,
+                             &QCoreApplication::quit);
+            t.start();
+            // wait
+            app.exec();
+        }
+    }
     else if (parser.isSet(configArgument)) { // CONFIG
         bool autostart = parser.isSet(autostartOption);
         bool filename = parser.isSet(filenameOption);
diff --git a/src/utils/screengrabber.cpp b/src/utils/screengrabber.cpp
index 58977b1c..f7d34305 100644
--- a/src/utils/screengrabber.cpp
+++ b/src/utils/screengrabber.cpp
@@ -88,3 +88,31 @@ QPixmap ScreenGrabber::grabEntireDesktop(bool &ok) {
     p.setDevicePixelRatio(QApplication::desktop()->devicePixelRatio());
     return p;
 }
+
+QPixmap ScreenGrabber::grabScreen(int screenNumber, bool &ok) {
+    QPixmap p;
+    bool isVirtual = QApplication::desktop()->isVirtualDesktop();
+    if (isVirtual || m_info.waylandDectected()) {
+        p = grabEntireDesktop(ok);
+        if (ok) {
+            QPoint topLeft(0, 0);
+#ifdef Q_OS_WIN
+            for (QScreen *const screen : QGuiApplication::screens()) {
+                QPoint topLeftScreen = screen->geometry().topLeft();
+                if (topLeft.x() > topLeftScreen.x() ||
+                        topLeft.y() > topLeftScreen.y()) {
+                    topLeft = topLeftScreen;
+                }
+            }
+#endif
+            QRect geometry = QApplication::desktop()->
+                    screenGeometry(screenNumber);
+            geometry.moveTo(geometry.topLeft() - topLeft);
+            p = p.copy(geometry);
+        }
+    } else {
+        p = QApplication::desktop()->screen(screenNumber)->grab();
+        ok = true;
+    }
+    return p;
+}
diff --git a/src/utils/screengrabber.h b/src/utils/screengrabber.h
index 6decab09..8ddef220 100644
--- a/src/utils/screengrabber.h
+++ b/src/utils/screengrabber.h
@@ -25,6 +25,7 @@ class ScreenGrabber : public QObject {
 public:
     explicit ScreenGrabber(QObject *parent = nullptr);
     QPixmap grabEntireDesktop(bool &ok);
+    QPixmap grabScreen(int screenNumber, bool &ok);
 
 private:
     DesktopInfo m_info;