-
Notifications
You must be signed in to change notification settings - Fork 2
/
screenshot.py
151 lines (118 loc) · 4.38 KB
/
screenshot.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
import ida_idaapi
import ida_kernwin
from PyQt5.QtCore import QPoint, QRect, QSize
from PyQt5.QtGui import QPixmap, QRegion
from PyQt5.QtWidgets import QApplication, QWidget
def scale_size(r: QRect, scale: int) -> QSize:
s = QRect(r)
s.setWidth(r.width() * scale)
s.setHeight(r.height() * scale)
return s.size()
def render_widget_img(widget: QWidget, scale: int) -> QPixmap:
rect = widget.rect()
img = QPixmap(scale_size(rect, scale))
img.setDevicePixelRatio(scale)
widget.render(img, QPoint(), QRegion(rect))
return img
class screenshot_handler_t(ida_kernwin.action_handler_t):
widget_only: bool
save_to_file: bool
def __init__(self, widget_only: bool, save_to_file: bool):
ida_kernwin.action_handler_t.__init__(self)
self.widget_only = widget_only
self.save_to_file = save_to_file
def activate(self, ctx): # pyright: ignore
if self.widget_only:
cur_twidget = ida_kernwin.get_current_widget()
target = ida_kernwin.PluginForm.TWidgetToPyQtWidget(cur_twidget)
else:
target = QApplication.activeWindow()
if target is None:
print("Could not find widget or window!")
return 1
scale = ida_kernwin.ask_long(2, "Screenshot scale multiplier:")
if not scale or scale < 1:
scale = 1
img = render_widget_img(target, scale)
if self.save_to_file:
path = ida_kernwin.ask_file(True, "screenshot.png", "Save screenshot")
if path:
img.save(path)
else:
QApplication.clipboard().setImage(img.toImage())
return 1
def update(self, ctx):
return ida_kernwin.AST_ENABLE_ALWAYS
ACTION_CAPTURE_WIDGET_COPY = "screenshot:CaptureWidgetCopy"
ACTION_CAPTURE_WINDOW_COPY = "screenshot:CaptureWindowCopy"
ACTION_CAPTURE_WIDGET_SAVE = "screenshot:CaptureWidgetSave"
ACTION_CAPTURE_WINDOW_SAVE = "screenshot:CaptureWindowSave"
class screenshot_ui_hooks_t(ida_kernwin.UI_Hooks):
popup_actions = []
def finish_populating_widget_popup(self, widget, popup): # pyright: ignore
for action in [
ACTION_CAPTURE_WIDGET_COPY,
ACTION_CAPTURE_WINDOW_COPY,
ACTION_CAPTURE_WIDGET_SAVE,
ACTION_CAPTURE_WINDOW_SAVE,
]:
ida_kernwin.attach_action_to_popup(
widget,
popup,
action,
f"Sc&reenshot/",
ida_kernwin.SETMENU_APP,
)
class screenshot_plugin_t(ida_idaapi.plugin_t):
flags = ida_idaapi.PLUGIN_DRAW | ida_idaapi.PLUGIN_HIDE
help = ""
comment = "Screenshot Capture"
wanted_name = "screenshot"
wanted_hotkey = ""
ui_hooks: screenshot_ui_hooks_t
def init(self):
self.ui_hooks = screenshot_ui_hooks_t()
ida_kernwin.register_action(
ida_kernwin.action_desc_t(
ACTION_CAPTURE_WIDGET_COPY,
"Copy wid~g~et screenshot to clipboard",
screenshot_handler_t(True, False),
None,
"Copy a screenshot of the current widget to the clipboard",
)
)
ida_kernwin.register_action(
ida_kernwin.action_desc_t(
ACTION_CAPTURE_WINDOW_COPY,
"Copy ~w~indow screenshot to clipboard",
screenshot_handler_t(False, False),
None,
"Copy a screenshot of the current window to the clipboard",
)
)
ida_kernwin.register_action(
ida_kernwin.action_desc_t(
ACTION_CAPTURE_WIDGET_SAVE,
"Save widget screenshot to file...",
screenshot_handler_t(True, True),
None,
"Save a screenshot of the current widget to a file",
)
)
ida_kernwin.register_action(
ida_kernwin.action_desc_t(
ACTION_CAPTURE_WINDOW_SAVE,
"Save window screenshot to file...",
screenshot_handler_t(False, True),
None,
"Save a screenshot of the current window to a file",
)
)
self.ui_hooks.hook()
return ida_idaapi.PLUGIN_KEEP
def run(self): # pyright: ignore
pass
def term(self):
self.ui_hooks.unhook()
def PLUGIN_ENTRY():
return screenshot_plugin_t()