KonaChanWallpaper/main.py

#!/usr/bin/env python3

# Copyright (c) 2015 Tom Li
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.


import os
import sys

import re
import json
from subprocess import call
from time import sleep
import argparse

import pycurl
import curl


# TODO:
# * tag filters
# * better error handling


def fequal(a, b):
    return abs(a - b) < 0.001


def ratio(a, b):
    return a / b


def detect_desktop_environment():
    desktop_env = "unknown"
    if os.environ.get("KDE_FULL_SESSION") == "true":
        desktop_env = "kde"
    elif os.environ.get("GNOME_DESKTOP_SESSION_ID"):
        desktop_env = "gnome"
    elif os.environ.get("MATE_DESKTOP_SESSION_ID"):
        desktop_env = "mate"
    elif sys.platform == "win32":
        desktop_env = "windows"
    return desktop_env


def set_curl_options(curl):
    try:
        curl.set_option(pycurl.HTTP_VERSION, pycurl.CURL_HTTP_VERSION_2_0)
    except AttributeError:
        pass

    ssl_library = pycurl.version_info()[5]
    if "OpenSSL" in ssl_library or "LibreSSL" in ssl_library:
        curl.set_option(pycurl.SSL_CIPHER_LIST, "ECDHE-ECDSA-AES256-GCM-SHA384:ECDHE-RSA-AES256-GCM-SHA384")
    elif "GnuTLS".lower() in ssl_library.lower():  # not sure about the capitalization, use lower case
        curl.set_option(pycurl.SSL_CIPHER_LIST, "PFS")
    else:
        print("Warning: unknown/untested SSL/TLS library")


def pick_up_a_url(curl, r18=False):
    URL_SAFE = "https://konachan.net/post/random"
    URL_R18 = "https://konachan.com/post/random"

    # XXX: Side-effect on the curl instance, but it doesn't matter.
    curl.set_option(pycurl.FOLLOWLOCATION, False)

    if r18:
        curl.get(URL_R18)
    else:
        curl.get(URL_SAFE)

    url = curl.get_info(pycurl.REDIRECT_URL)
    print("Selected", url)
    return url


def fetch_image_info(curl, image_page_url):
    html = curl.get(image_page_url).decode("UTF-8")
    json_info = re.findall('(?<=Post.register_resp\\().*(?=\\);)', html)[0]
    return json.loads(json_info)


def download_image(curl, image_url, filename):

    def progress(dwn_total, dwn, up_total, up):
        global progress_prev
        if dwn_total == 0:
            return
        progress = int(dwn / dwn_total * 100)
        if progress != progress_prev and progress % 10 == 0:
            print("downloaded", "%d%%" % progress)
        progress_prev = progress

    print("\nnow downloading %s..." % filename)

    # XXX: Side-effect on the curl instance, but it doesn't matter.
    global progress_prev
    progress_prev = 0
    curl.set_option(pycurl.NOPROGRESS, 0)
    curl.set_option(pycurl.PROGRESSFUNCTION, progress)

    # HACK: curl doesn't truncate the old data, multiple GETs would not flush the old
    # data from the buffer and it results in returning both previous data and new data.
    # It is uncertain if this behavior is a feature or a bug due to the fact that curl.Curl()
    # is rarely used and poorly documented. After all, it is just a simple class around 100
    # lines... Thus we manually truncate the data here.
    # See also: https://stackoverflow.com/questions/4330812/how-do-i-clear-a-stringio-object
    curl.payload_io.truncate(0)
    curl.payload_io.seek(0)

    data = curl.get(image_url)
    with open(filename, "wb") as file:
        file.write(data)


def set_wallpaper(path):
    path = os.path.realpath(path)
    desktop = detect_desktop_environment()
    if desktop == "gnome":
        path = "file://%s" % path
        call(("gsettings", "set", "org.gnome.desktop.background", "picture-uri", path))
    elif desktop == "mate":
        call(("gsettings", "set", "org.mate.background", "picture-filename", path))
    elif desktop == "kde":
        raise NotImplementedError("How to change wallpaper for KDE?")
    elif desktop == "windows":
        path = '"%s"' % path
        call(("reg", "add", "HKCU\Control Panel\Desktop", "/v", "Wallpaper", "/f", "/t", "REG_SZ", "/d", path))
    else:
        raise NotImplementedError("I don't know how to do it")


def main():
    parser = argparse.ArgumentParser(description="Fetch a random wallpaper from Konachan")
    parser.add_argument("--width", dest="width", type=int, default=0,
                        help="download only if the ratio of the image not matches with your suggestion")
    parser.add_argument("--height", dest="height", type=int, default=0,
                        help="download only if the ratio of the image not matches with your suggestion")
    parser.add_argument("--r18", dest="r18", action="store_true", default=False,
                        help="allow to fetch explicit (R-18) images")
    args = parser.parse_args()

    c = curl.Curl()
    set_curl_options(c)

    while 1:
        random_url = pick_up_a_url(c, args.r18)
        full_info = fetch_image_info(c, random_url)

        image_info = full_info["posts"][0]
        if (args.height and args.width and
           not fequal(ratio(image_info["height"], image_info["width"]), ratio(args.height, args.width))):
            print("ratio is not perfect, try again...")
            sleep(5)
        else:
            break

    print("tags:")
    tag_info = full_info["tags"]
    for key in tag_info.keys():
        print(key)

    filename = str(image_info["id"])
    if image_info["file_url"].startswith("https"):
        file_url = image_info["file_url"]
    else:
        file_url = "https:%s" % image_info["file_url"]

    download_image(c, file_url, filename)
    set_wallpaper(filename)


if __name__ == "__main__":
    main()