#!/usr/bin/env python3


import socket
import os
import struct
import time


sockr_path = '/etc/lxc/rebinds/devlog/log'


def get_namespace_from_proc(pid):
    '''Get lxc cgroup namespace name associated with pid, or None'''
    try:
        with open('/proc/%s/cgroup' % pid, 'rb') as fin:
            data = fin.read()
        data = data.splitlines()[0]
        truncation_index = data.index(b'/lxc/')
        return data[truncation_index+len(b'/lxc/'):]
    except (FileNotFoundError, ProcessLookupError, ValueError): pass


def get_namespace(pid, namespace_cache=dict(), timeout_time=30):
    '''
        Get cached namespace name (works only with /lxc/* cgroup names). If
        request is made on a name which is cached, the timeout of the name is
        reset. If the name is not found or the timeout is over, the name is
        got from proc and cached.

        There is a catch though: if there is pid overflow on the host system
        (pid restart from count 1) there might be a chance the pid is no
        longer on the same cached namespace name. But the chance are very
        unlikely and the consequences are ok for my usage. If this is
        important for you, you can set timeout_time to 0 (or rewrite this
                part). Thus, each time a line is received from /dev/log in
        a container, an open call is done to get the name which is bad for
        performances.
    '''
    tcur = time.time()
    name, timeout = namespace_cache.get(pid, (None, None))
    if name and timeout >= tcur:
        namespace_cache[pid][1] = tcur + timeout_time
        #print('debug: got cached name')
        return name
    else:
        name = get_namespace_from_proc(pid)
        timeout = tcur + timeout_time
        namespace_cache[pid] = [name, timeout]
        #print('debug: name cache miss')
    return name


def main():
    sockr = socket.socket(socket.AF_UNIX, socket.SOCK_DGRAM)
    sockr.setsockopt(socket.SOL_SOCKET, socket.SO_PASSCRED, 1)
    os.unlink(sockr_path)
    sockr.bind(sockr_path)
    os.chmod(sockr_path, mode=0o666)
    socks = socket.socket(socket.AF_UNIX, socket.SOCK_DGRAM)
    while True:
        data = sockr.recvmsg(8192, 256)

        pid = 0
        for cmsg in data[1]:
            if cmsg[0] == socket.SOL_SOCKET and cmsg[1] == socket.SCM_CREDENTIALS:
                pid = struct.unpack('i', cmsg[2][0:4])[0]

        data_mangled = data[0]
        name = pid and get_namespace(pid)
        if name:
            first_space_index = data[0].find(b' ')
            marker_end_index = data[0].find(b': ', first_space_index + 16)
            insert_index = data[0].rfind(b' ', 0, marker_end_index) + 1
            data_mangled = data[0][0:insert_index] + name + b'/' + data[0][insert_index:]

        #print('debug:', data_mangled)
        socks.sendto(data_mangled, '/dev/log')


if __name__ == '__main__':
    main()