#! /bin/python3.6
# -*- coding: utf-8 -*-
#
# This file is part of INGInious. See the LICENSE and the COPYRIGHTS files for
# more information about the licensing of this file.
import os
import shlex
import socket
import subprocess
import threading

import msgpack
import array

import resource


def handle_signals(subprocess, socket):
    """ Handles signals given by run_student on the socket """
    while True:
        signal = socket.recv(3)
        if signal == b'---' or len(signal) < 3: # quit
            return
        print("received %s" % signal.decode("utf8"))
        subprocess.send_signal(int(signal.decode('utf8')))

def recv_fds(sock, msglen, maxfds):
    """ Receive FDs from the unix socket. Copy-pasted from the Python doc """
    fds = array.array("i")  # Array of ints
    msg, ancdata, flags, addr = sock.recvmsg(msglen, socket.CMSG_LEN(maxfds * fds.itemsize))
    for cmsg_level, cmsg_type, cmsg_data in ancdata:
        if (cmsg_level == socket.SOL_SOCKET and cmsg_type == socket.SCM_RIGHTS):
            # Append data, ignoring any truncated integers at the end.
            fds.fromstring(cmsg_data[:len(cmsg_data) - (len(cmsg_data) % fds.itemsize)])
    return msg, list(fds)

def setlimits():
    os.setgid(4242)
    os.setuid(4242)
    resource.setrlimit(resource.RLIMIT_NPROC, (1000, 1000))

# Connect to the socket
client = socket.socket(socket.AF_UNIX)  # , socket.SOCK_CLOEXEC) # for linux only
client.connect("/__parent.sock")

# Say hello
print("Saying hello")
client.send(b'H')
print("Said hello")

# Receive fds
print("Receiving fds")
msg, fds = recv_fds(client, 1, 3)
assert msg == b'S'
print("Received fds")

# Unpack the start message # TODO: it's ugly
print("Unpacking start cmd")
unpacker = msgpack.Unpacker()
start_cmd = None
while start_cmd is None:
    s = client.recv(1)
    unpacker.feed(s)
    for obj in unpacker:
        start_cmd = obj
print("Unpacked start cmd")

# Add some elements to /etc/hosts and /etc/resolv.conf if needed
system_files = {"hosts": ("/etc/hosts", True), "resolv.conf": ("/etc/resolv.conf", False)}
for name, (spath, append) in system_files.items():
    if os.path.exists(os.path.join('/task/systemfiles/', name)):
        try:
            open(spath, 'ab' if append else 'wb').write(b'\n' + open(os.path.join('/task/systemfiles/', name), 'rb').read())
        except IOError:
            exit(254)

# Start the process
print("Chdir")
os.chdir(start_cmd["working_dir"])
print("Popen")
print(fds)
p = subprocess.Popen(shlex.split(start_cmd["command"]), preexec_fn=setlimits, stdin=fds[0], stdout=fds[1], stderr=fds[2])

# Handle the signals
print("Signal thread")
thread = threading.Thread(target=lambda: handle_signals(p, client), daemon=True)
thread.start()

# Wait until the subprocess exit
print("Waiting for process to end")
retval = p.wait()

# Exit with the same return value than the process (it will be given to docker and given back to run_student)
print("Exiting")
exit(retval)
