screenshot.py

#!/usr/bin/env python3

"""Grabbing screenshots and transfer using FTP

Initial command: pip install wxpython
    * Not compatiable with python3.10, fix is use python3.9
    * Problematic, switched to PIL

FTP command: python -m pyftpdlib --verbose -d ftp -u admin -P admin -w

Usage: $ python3 screenshot.py <filename>
"""

import ftplib
import os
import socket
import sys

from PIL import ImageGrab


def transfer(screenshot: str) -> None:
    """Transfers screenshot over FTP

    Args:
        screenshot (str): Filename of the screenshot
    """
    if not os.path.isdir("ftp"):
        os.makedirs("ftp")
    try:
        ipaddr = socket.gethostbyname(socket.gethostname())
        session = ftplib.FTP(ipaddr, "admin", "admin")
        file = open(screenshot, "rb")
        session.storbinary(f"STOR {screenshot}", file)
        file.close()
        session.quit()
    except Exception as e:
        print(f"Error: {e}")
        print("FTP transfer failed, exiting...")
        sys.exit(1)


def screenshot(filename: str) -> str:
    """Takes a screenshot in the current context

    Args:
        filename (str): The name of the screenshot to be saved to

    Returns:
        str: The filename if successful else None
    """
    try:
        # Set up screenshot and save it
        # Source: https://pillow.readthedocs.io/en/stable/reference/ImageGrab.html
        screenshot = ImageGrab.grab()
        screenshot.save(f"{filename}")
        return filename
    except Exception as e:
        print(f"Error: {e}")
        print("Screenshot failed, exiting...")
        sys.exit(1)


if __name__ == "__main__":
    if len(sys.argv) < 2:
        print(f"Usage: python3 {sys.argv[0]} <filename>")
        sys.exit(1)
    result = screenshot(sys.argv[1])
    transfer(result)