#!/usr/bin/env python
# -*- coding:utf-8 -*-
"""
Удаленное выполнение команд на dupload.dist.yandex.ru от имени спец-пользователя robot-direct-dmover

для выполнения настоящих действий нужно sudo, dry run можно и без него

Формат запуска:

    dt-dist [команда] [название пакета/файла с пакетами]

    Список комманд:
    dt-dist --list

    Пробный запуск:
    DRY_RUN=1 direct-dist.py dmove_stable yandex-direct=1.82616-1

    Перемещение пакетов (dmove):
    dt-dist dmove_stable -r dt yandex-direct=1.82616-1              # перемещение из testing
    dt-dist dmove_testing_prestable -r dp yandex-direct=1.82616-1
    dt-dist dmove_prestable_stable -r dp yandex-direct=1.82616-1
    dt-dist dmove -r direct-trusty -f unstable -t testing yandex-direct=1.82616-1

    Массовое перемещение пакетов (dmove):
    dt-dist dmove_from_file -r dp -f unstable -t testing ~/filename
    dt-dist dmove_from_file -r dp -f unstable -t testing -          # из stdin
    dt-dist dmove_testing_from_file -r dp ~/filename
    dt-dist dmove_stable_from_file -r dp -                          # перемещение из testing
    dt-dist dmove_prestable_stable_from_file -r dp -

    Удаление пакетов (rmove)
    dt-dist rmove -r dmp -b unstable direct-moderate=1.11423.11493-1

    Копирование пакетов из другого репозитория
    dt-dist bmove -f dt -t dmp -b unstable yandex-du-ts-updater=1.60-1

    Поиск пакетов:
    dt-dist find_package yandex-direct_1.80838

Опции:
        -h, --help 
            показать справку и завершиться
        -l, --list
            вывести список команд

        -r <repository> 
            название репозитория (например, direct-trusty)
        -f, --from <branch/repo>
            название ветки или репозитория "откуда"
        -t, --to <branch/repo>
            название ветки или репозитория "куда"
        -b, --branch <branch>
            название ветки для команд, которым нужна только одна ветка (например, rmove и bmove)
        --no-ssh
            запустить dmove асинхронно (поддерживается не всеми командами)

"""

#    TODO
#    + rmove
#    + массовый dmove
#    - когда понадобится -- массовый rmove аналогично массовому dmove
#    - научиться безопасно развозить ключи по всем нужны машинам (ppcdev*)

import sys, re, os
import subprocess
import argparse
import time
import json
import requests


allowed_repos = [
    "direct-mod-precise", 
    "direct-precise", 
    "direct-trusty", 
    "direct-xenial", 
    "direct-bionic", 
    "direct-common"
    ]

aliases = {
        "dc": "direct-common",
        "dp": "direct-precise",
        "dt": "direct-trusty",
        "dx": "direct-xenial",
        "db": "direct-bionic",
        "dmp": "direct-mod-precise",
        }


def my_run(*args,**kwargs):
    if os.getenv('DRY_RUN', '') == "1":
        print "dry run: %s" % (" ".join(args))
    else:
        if kwargs.get("no_sudo", ""):
            my_args = ["ssh", "-o", "StrictHostKeyChecking=no", "dupload.dist.yandex.ru"]
        else:
            my_args = ["ssh",
                       "-l", "robot-direct-dmover",
                       "-i", "/etc/robot-direct-dmover-keys/id_rsa",
                       "-o", "StrictHostKeyChecking=no",
                       "dupload.dist.yandex.ru"]

        subprocess.call(my_args + list(args))


def normalize_repo(repo):
    if repo in aliases:
        return aliases[repo]
    else: 
        return repo


def _mass_dmove(repo, src, dst, to_dmove, no_ssh=False):
    """
    переместить много пакетов между репозиториями (unstable/testing/unstable)
    """
    norm_repo = normalize_repo(repo)
    if not norm_repo in allowed_repos:
        print "repo %s isn't allowed, stop" % (repo)
        exit(57)

    i = 0
    for task in to_dmove:
        i += 1
        if len(task) != 2:
            print "task #%s: wrong format '%s', stop" % (i, task)
            exit(57)

        if not re.match( r'^[a-z0-9-.]+$', task[0],):
            print "task #%s: incorrect package name '%s', stop\n" % (i, task[0])
            exit(57)

        if not re.match( r'^[a-z0-9-.~:]+$', task[1], re.I):
            print "task #%s: incorrect version '%s', stop\n" % (i, task[1])
            exit(57)

    if no_ssh:
        if os.environ.get("DRY_RUN", "") == "1":
            if to_dmove:
                print "dry run: ..."
            sys.exit(0)

        if norm_repo == "direct-trusty" and src == "unstable" and dst == "testing":
            for task in to_dmove:
                file_task = "/var/lib/dt-async-dmove/%s_%s_%s_%s_%s" % (norm_repo, src, dst, task[0], task[1])
                with open(file_task, "w") as f:
                    print "task %s created" % file_task 

            retries = 10
            for task in to_dmove:
                print "waiting task %s_%s_%s_%s_%s" % (norm_repo, src, dst, task[0], task[1])
                success = False
                for i in xrange(retries):
                    try:
                        output = requests.get("http://dist.yandex.ru/api/v1/search?pkg=%s&ver=%s&repo=%s&strict=true" % (task[0], task[1], norm_repo)).json()
                        if output['success'] and output['result'][0]['environment'] == 'testing':
                            success = True
                            break

                    except Exception as e:
                        print e

                    if i + 1 != retries:
                        time.sleep(10)

                if not success:
                    print "can't dmove %s=%s asynchronously" % (task[0], task[1])
                    exit(57)
                else:
                    print "successfully dmove'd"

        else:
            print "can't dmove without ssh with such parametres: repo=%s, src=%s, dst=%s" % (norm_repo, src, dst)
            exit(57)
    else:
        for task in to_dmove:
            my_run("sudo dmove %s %s %s %s %s" % (norm_repo, dst, task[0], task[1], src))


def dmove_from_file(repo, src, dst, filename, no_ssh=False):
    """
    переместить много пакетов между репозиториями (unstable/testing/unstable)
    пакеты и версии читаются из файла, по одному на строку: yandex-direct[ ,=_]1.23456-1
    если вместо имени файла '-' (дефис) -- читается stdin
    """
    if filename == "-": 
         fd = sys.stdin
    else: 
         fd = open(filename, 'r')

    # зачитываем весь файл, парсим, ничего не валидируем; валидация -- в _mass_dmove
    to_dmove = []
    for line in fd:
        line = line.rstrip()
        if re.match( r'^#', line):
            continue
        if re.match( r'^\s*$', line):
            continue

        parts = re.split("[ ,=_]", line)

        to_dmove += [ parts ]

    _mass_dmove(repo, src, dst, to_dmove, no_ssh)


def dmove(repo, src, dst, package, version, no_ssh=False):
    """
    переместить пакет между репозиториями (unstable/testing/unstable)
    """
    _mass_dmove(repo, src, dst, [ [package, version] ], no_ssh)


def dmove_testing(repo, package, version, no_ssh):
    """
    переместить пакет из unstable в testing
    """
    dmove(repo, "unstable", "testing", package, version, no_ssh)


def dmove_prestable(repo, package, version):
    """
    переместить пакет из unstable в prestable
    """
    dmove(repo, "unstable", "prestable", package, version)


def dmove_testing_prestable(repo, package, version):
    """
    переместить пакет из testing в prestable
    """
    dmove(repo, "testing", "prestable", package, version)


def dmove_stable(repo, package, version):
    """
    переместить пакет из testing в stable
    """
    dmove(repo, "testing", "stable", package, version)


def dmove_prestable_stable(repo, package, version):
    """
    переместить пакет из prestable в stable
    """
    dmove(repo, "prestable", "stable", package, version)


def dmove_testing_from_file(repo, filename, no_ssh=False):
    """
    переместить много пакетов из unstable в testing; список пакетов берется из файла; "-" -- stdin
    """
    dmove_from_file(repo, "unstable", "testing", filename, no_ssh)


def dmove_prestable_from_file(repo, filename):
    """
    переместить много пакетов из unstable в prestable; список пакетов берется из файла; "-" -- stdin
    """
    dmove_from_file(repo, "unstable", "prestable", filename)


def dmove_testing_prestable_from_file(repo, filename):
    """
    переместить много пакетов из testing в prestable; список пакетов берется из файла; "-" -- stdin
    """
    dmove_from_file(repo, "testing", "prestable", filename)


def dmove_stable_from_file(repo, filename):
    """
    переместить много пакетов из testing в stable; список пакетов берется из файла; "-" -- stdin
    """
    dmove_from_file(repo, "testing", "stable", filename)


def dmove_prestable_stable_from_file(repo, filename):
    """
    переместить много пакетов из prestable в stable; список пакетов берется из файла; "-" -- stdin
    """
    dmove_from_file(repo, "prestable", "stable", filename)


def find_package(package, version=''):
    """
    поиск пакета на dupload.dist.yandex.ru
    """
    my_run('find_package.sh %s%s' % (package, ("_" + version) if version else ""), no_sudo=True)


#rmove <branch> <repo> <package> <version>
def rmove(repo, branch, package, version):
    """
    удалить пакет
    """
    norm_repo = normalize_repo(repo)
    if not norm_repo in allowed_repos:
        sys.exit("repo %s isn't allowed, stop" % (repo))
    my_run("sudo rmove %s %s %s %s" % (norm_repo, branch, package, version))


# bmove <-c|-m> <branch> <torepo> <package> <version> <fromrepo>
def bmove(src, dst, branch, package, version):
    """
    скопировать пакет в репозиторий
    """
    fromrepo = normalize_repo(src)
    torepo = normalize_repo(dst)
    if not torepo in allowed_repos:
        sys.exit("copying to repo %s is not allowed, stop" % (torepo))
    # у робота есть пишущий доступ ко всем репозиториям на dist'е, поэтому разрешено только копирование (-c), но не перемещение (-m), чтобы не удалить случайно пакет из "чужого" репозитория
    # для того, чтобы после копирования удалить пакет из "нашего" репозитория-источника, есть rmove
    my_run("sudo bmove -c %s %s %s %s %s" % (branch, torepo, package, version, fromrepo))


def parse_options():
    global opts

    parser = argparse.ArgumentParser(description="Работа с dist'ом", add_help=False, usage=argparse.SUPPRESS)

    parser.add_argument("-h", "--help", dest="help", action='store_true', help="справка", default=None)
    parser.add_argument("-l", "--list", dest="list", help="вывести список команд", action='store_true', default=None)

    opts, extra = parser.parse_known_args()

    if opts.help:
        print __doc__
        sys.exit(0)

    if opts.list:
        print "\n".join(cmds.keys())
        sys.exit(0)

    parser.add_argument("command", type=str, default=None, help="команда для запуска")
    parser.add_argument("package", type=str, default=None, help="название пакета")

    parser.add_argument("-r", dest="repo", help="название репозитория", type=str, default=None, action='append')
    parser.add_argument("-f", "--from", dest="src", help="ветка/репозиторий from", type=str, default=None)
    parser.add_argument("-t", "--to", dest="dst", help="ветка/репозиторий to", type=str, default=None)
    parser.add_argument("-b", "--branch", dest="branch", help="название ветки", type=str, default=None)
    parser.add_argument("--no-ssh", dest="no_ssh", help="запустить dmove асинхронно", action='store_true', default=None)

    opts, extra = parser.parse_known_args()

    if len(extra) > 0:
        sys.exit("There are unknown parameters")

    if opts.repo:
        if len(opts.repo) > 1:
            sys.exit("you must specify no more than one repository")
        else:
            opts.repo = opts.repo[0]

    if opts.package:
        splitted = re.split(ur'[,=_]', opts.package, maxsplit=1)
        if len(splitted) > 1:
            opts.package, opts.version = splitted

    if opts.command not in cmds:
        for main_cmd in cmd_aliases:
            if opts.command in cmd_aliases[main_cmd]:
                opts.command = main_cmd
                break

    # пока поддерживается только dmove в testing, по дефолту выставляем параметр no-ssh
    # потом возможно надо переделать логику с этим параметром
    if opts.command in ["dmove_testing", "dmove_testing_from_file"] or opts.command in ["dmove", "dmove_from_file"] and opts.dst == "testing":
        norm_repo = normalize_repo(opts.repo)
        if norm_repo == "direct-trusty":
            opts.no_ssh = True

    return opts


cmd_aliases = {
    'find_package': ['find', 'find-package'],
    'dmove_testing': ['dmove-testing'],
    'dmove_stable': ['dmove-stable'],
}


cmds = {
    'dmove': {
        'fn': dmove,
        'args': ['repo', 'src', 'dst', 'package', 'version'],
        'extra_args': ['no_ssh']
    },
    'dmove_testing': {
        'fn': dmove_testing,
        'args': ['repo', 'package', 'version'],
        'extra_args': ['no_ssh']
    },
    'dmove_prestable': {
        'fn': dmove_prestable,
        'args': ['repo', 'package', 'version']
    },
    'dmove_testing_prestable': {
        'fn': dmove_testing_prestable,
        'args': ['repo', 'package', 'version']
    },
    'dmove_stable': {
        'fn': dmove_stable,
        'args': ['repo', 'package', 'version']
    },
    'dmove_prestable_stable': {
        'fn': dmove_prestable_stable,
        'args': ['repo', 'package', 'version']
    },
    'dmove_from_file': {
        'fn': dmove_from_file,
        'args': ['repo', 'src', 'dst', 'package'],
        'extra_args': ['no_ssh']
    },
    'dmove_testing_from_file': {
        'fn': dmove_testing_from_file,
        'args': ['repo', 'package'],
        'extra_args': ['no_ssh']
    },
    'dmove_prestable_from_file': {
        'fn': dmove_prestable_from_file,
        'args': ['repo', 'package']
    },
    'dmove_testing_prestable_from_file': {
        'fn': dmove_testing_prestable_from_file,
        'args': ['repo', 'package']
    },
    'dmove_stable_from_file': {
        'fn': dmove_stable_from_file,
        'args': ['repo', 'package']
    },
    'dmove_prestable_stable_from_file': {
        'fn': dmove_prestable_stable_from_file,
        'args': ['repo', 'package']
    },
    'rmove': {
        'fn': rmove,
        'args': ['repo', 'branch', 'package', 'version']
    },
    'bmove': {
        'fn': bmove,
        'args': ['src', 'dst', 'branch', 'package', 'version']
    },
    'find_package': {
        'fn': find_package,
        'args': ['package'],
        'extra_args': ['version'],
    },
}


def run():
    opts = parse_options()
    opts_values = vars(opts)
    args = []

    for arg in cmds[opts.command]['args']:
        if arg not in opts_values or opts_values[arg] is None:
            sys.exit("error: parameter %s is required" % arg)

        args.append(opts_values[arg])

    if 'extra_args' in cmds[opts.command]:
        for arg in cmds[opts.command]['extra_args']:
            if arg in opts_values:
                args.append(opts_values[arg])
            else:
                args.append(None)

    for arg in opts_values:
        if arg != 'command' and opts_values[arg] is not None and arg not in cmds[opts.command]['args'] and ('extra_args' not in cmds[opts.command] or arg not in cmds[opts.command]['extra_args']):
            sys.exit("error: parameter %s is not for %s" % (arg, opts.command))

    cmds[opts.command]['fn'](*args)


if __name__ == '__main__':
    run()

