#!/usr/bin/env python3

# Copyright (C) 2018 Simon Quigley <tsimonq2@ubuntu.com>
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program. If not, see <http://www.gnu.org/licenses/>.

import json
import libcalamares
from time import strftime
import urllib.request
from urllib.error import HTTPError
from urllib.error import URLError
import socket
import logging
from lsb_release import get_distro_information

global sources
sources = """# Automatically generated by Calamares on DATE.
# Lines starting with "deb" are mandatory, while lines starting with "deb-src"
# are for more detailed package information.

## See http://help.ubuntu.com/community/UpgradeNotes for how to upgrade to
## newer versions of DISTRIBUTION.
deb URL/ubuntu/ CODENAME main restricted
# deb-src URL/ubuntu/ CODENAME main restricted

## Major bug fix updates produced after the final release of DISTRIBUTION.
## Have you noticed a regression? Please report it!
## https://wiki.ubuntu.com/StableReleaseUpdates#Regressions
deb URL/ubuntu/ CODENAME-updates main restricted
# deb-src URL/ubuntu/ CODENAME-updates main restricted

## Software from this repository is ENTIRELY UNSUPPORTED by the Ubuntu team.
## Also, please note that software in Universe WILL NOT receive any review or
## updates from the Ubuntu security team directly. Updates in this repository
## are provided by volunteers, but most come from Debian.
deb URL/ubuntu/ CODENAME universe
# deb-src URL/ubuntu/ CODENAME universe
deb URL/ubuntu/ CODENAME-updates universe
# deb-src URL/ubuntu/ CODENAME-updates universe

## Software from this repository is ENTIRELY UNSUPPORTED by the Ubuntu team,
## and may not be under a free licence. Please satisfy yourself as your rights
## to use the software. Also, please note that software in Multiverse WILL NOT
## receive any review or updates from the Ubuntu security team directly.
deb URL/ubuntu/ CODENAME multiverse
# deb-src URL/ubuntu/ CODENAME multiverse
deb URL/ubuntu/ CODENAME-updates multiverse
# deb-src URL/ubuntu/ CODENAME-updates multiverse

## Software from this repository contains tested security updates from the
## Ubuntu security team.
deb http://security.ubuntu.com/ubuntu CODENAME-security main restricted
# deb-src http://security.ubuntu.com/ubuntu CODENAME-security main restricted
deb http://security.ubuntu.com/ubuntu CODENAME-security universe
# deb-src http://security.ubuntu.com/ubuntu CODENAME-security universe
deb http://security.ubuntu.com/ubuntu CODENAME-security multiverse
# deb-src http://security.ubuntu.com/ubuntu CODENAME-security multiverse

## Software from this repository may not have been tested as extensively as
## software contained in the main release, although it includes newer versions
## of some applications which may provide useful features. Also, please note
## that software in Backports WILL NOT receive any review or updates from the
## Ubuntu security team.
deb URL/ubuntu/ CODENAME-backports main restricted universe multiverse
# deb-src URL/ubuntu/ CODENAME-backports main restricted universe multiverse

## Uncomment the following two lines to add software from Canonical's
## "partner" repository.
## This software is not part of Ubuntu, but is offered by Canonical and the
## respective vendors as a service to Ubuntu users.
# deb http://archive.canonical.com/ubuntu CODENAME partner
# deb-src http://archive.canonical.com/ubuntu CODENAME partner"""

SUBDOMAINS_BY_COUNTRY_CODE = {"US": "us.",
                              "AU": "au.",
                              "SE": "no.",
                              "NO": "no.",
                              "NZ": "nz.",
                              "NL": "nl.",
                              "KR": "kr.",
                              "DE": "de.",
                              "GE": "ge.",
                              "PF": "pf.",
                              "CZ": "cz.",
                              "HR": "hr."}


def getcountrycode():
    """
    Return the two-letter country code or an empty string.

    Tries to determine the country code based on the public IP address,
    if the device is connected to the Internet. Otherwise it returns
    an empty string.
    """
    if libcalamares.globalstorage.value("hasInternet"):
        geoipurl = libcalamares.job.configuration["geoip"]["url"]
        try:
            with urllib.request.urlopen(geoipurl, timeout=75) as http_response:
                localedata = json.loads(http_response.read().decode())
        except HTTPError as http_error:
            logging.error("Data not retrieved because %s - URL: %s",
                          http_error, geoipurl)
        except URLError as url_error:
            if isinstance(url_error.reason, socket.timeout):
                logging.error("Socket timed out - URL %s", geoipurl)
            else:
                logging.error("Non-timeout protocol error.")
        else:
            logging.info("Country successfully determined.")
            return localedata["country"]
    else:
        return ""


def get_subdomain_by_country(countrycode):
    """Return the subdomain for the given countrycode
    or an empty string.
    """
    if countrycode in SUBDOMAINS_BY_COUNTRY_CODE.keys():
        return SUBDOMAINS_BY_COUNTRY_CODE[countrycode]
    else:
        return ""


def getcodename():
    """Return the codename of the distribution, similar to lsb_release -cs"""
    return get_distro_information()["CODENAME"]


def changesources(subdomain):
    """Replace the placeholders and then create the sources.list"""
    distro = libcalamares.job.configuration["distribution"]
    url = "http://{}{}".format(subdomain,
                               libcalamares.job.configuration["baseUrl"])

    global sources
    sources = sources.replace("DISTRIBUTION", distro)
    sources = sources.replace("CODENAME", getcodename())
    sources = sources.replace("URL", url)
    sources = sources.replace("DATE", strftime("%Y-%m-%d"))

    filepath = libcalamares.globalstorage.value("rootMountPoint")
    filepath += "/etc/apt/sources.list"
    with open(filepath, "r+") as sourcesfile:
        sourcesfile.seek(0)
        sourcesfile.write(sources)
        sourcesfile.truncate()


def run():
    """Autoselect a mirror and create the sources.list file."""
    countrycode = getcountrycode()
    subdomain = get_subdomain_by_country(countrycode)

    changesources(subdomain)
