# -*- coding: utf-8 -*-
from __future__ import print_function
import base64, os, socket, subprocess, tempfile, sys, binascii
import re
from datetime import datetime

# ---- Compatibility for CentOS 6 (Py2.6/2.7) ----
# Safe print (UTF-8, hỗ trợ stream=file)
def _p(msg, color=None, stream=None, **kwargs):
    # Back-compat cho calls kiểu file=...
    if stream is None and 'file' in kwargs:
        stream = kwargs['file']
    if stream is None:
        stream = sys.stdout
    try:
        unicode  # Py2 check
        b = msg.encode('utf-8') if isinstance(msg, unicode) else msg
    except NameError:
        b = msg
    try:
        if color:
            stream.write(color)
        stream.write(b)
        if color:
            try:
                stream.write(Fore.RESET)
            except Exception:
                pass
        stream.write("\n")
    except Exception:
        try:
            stream.write((u'' + msg).encode('utf-8') + "\n")
        except Exception:
            sys.stdout.write((u'' + msg).encode('utf-8') + "\n")

# optparse (có sẵn trên Py2.6)
try:
    from optparse import OptionParser
except Exception:
    OptionParser = None  # fallback rất cũ

# Color handling (tùy chọn)
try:
    from colorama import Fore, init as colorama_init
    colorama_init()
except Exception:
    class _NoColor(object):
        def __getattr__(self, name): return ''
    Fore = _NoColor()

BUNDLED_SCRIPT = "audit_cis_centos6.sh"


def resource_path(relative_path):
    if getattr(sys, "frozen", False):
        return os.path.join(sys._MEIPASS, relative_path)
    return os.path.join(os.path.dirname(os.path.abspath(__file__)), relative_path)


def resolve_script_path(user_path=None):
    if user_path and os.path.isfile(user_path):
        return os.path.abspath(user_path)
    bundled = resource_path(BUNDLED_SCRIPT)
    if os.path.isfile(bundled):
        return bundled
    same_dir = os.path.join(os.getcwd(), BUNDLED_SCRIPT)
    if os.path.isfile(same_dir):
        return same_dir
    return None


def detect_os_name():
    """
    Trả về chuỗi tên HĐH ngắn gọn.
    Ưu tiên /etc/*-release; fallback sang uname -srm.
    """
    paths = [
        '/etc/centos-release',
        '/etc/redhat-release',
        '/etc/os-release',      # có trên bản mới hơn
        '/etc/lsb-release',     # có thể chứa DISTRIB_DESCRIPTION=
    ]
    for p in paths:
        try:
            if os.path.isfile(p):
                data = open(p, 'rb').read()
                try:
                    text = data.decode('utf-8', 'ignore')
                except Exception:
                    text = data
                # /etc/lsb-release
                if 'DISTRIB_DESCRIPTION=' in text:
                    for line in text.splitlines():
                        if line.startswith('DISTRIB_DESCRIPTION='):
                            val = line.split('=',1)[1].strip().strip('"')
                            return val
                # các file *release khác: dùng dòng đầu
                first = text.splitlines()[0].strip()
                if first:
                    return first
        except Exception:
            pass
    # Fallback: uname
    try:
        out = subprocess.Popen("uname -srm", shell=True,
                               stdout=subprocess.PIPE, stderr=subprocess.PIPE).communicate()[0]
        try:
            return out.decode('utf-8','ignore').strip()
        except Exception:
            return out.strip()
    except Exception:
        return 'Unknown-OS'

try:
    hostname = socket.gethostname()
    IPAddr = socket.gethostbyname(hostname)
except Exception:
    hostname = 'unknown-host'
    IPAddr = '127.0.0.1'

cwd = os.getcwd()


def pkcs7_unpad(data, block_size=16):
    if not data:
        return data
    pad_len = ord(data[-1]) if isinstance(data[-1], str) else data[-1]
    if pad_len < 1 or pad_len > block_size:
        # Defensive: return as-is to avoid exception storms in Py2 envs
        return data
    return data[:-pad_len]


def write_content(content, use_oaep=True):
    """
    Hybrid AES-256-CBC + RSA encryption using OpenSSL CLI.
    Only 1 RSA operation for the AES key; bulk data encrypted with fast AES.
    Output: HYBRID_V1:<b64_enc_key>:<b64_iv>:<b64_ciphertext>
    """
    import binascii

    pubkey_path = os.path.join(os.getcwd(), 'public_key.pem')
    if not os.path.isfile(pubkey_path):
        raise RuntimeError("public_key.pem khong ton tai tai: %s" % pubkey_path)

    try:
        unicode
        if isinstance(content, unicode):
            message = content.encode('utf-8')
        else:
            message = content
    except NameError:
        if isinstance(content, str):
            message = content.encode('utf-8')
        else:
            message = content

    aes_key = os.urandom(32)
    iv = os.urandom(16)

    fd_in, path_in = tempfile.mkstemp(prefix='aes_in_', dir='/tmp')
    fd_out, path_out = tempfile.mkstemp(prefix='aes_out_', dir='/tmp')
    os.close(fd_in)
    os.close(fd_out)

    try:
        with open(path_in, 'wb') as f:
            f.write(message)

        key_hex = binascii.hexlify(aes_key)
        iv_hex = binascii.hexlify(iv)

        cmd_aes = [
            '/usr/bin/openssl', 'enc', '-aes-256-cbc',
            '-K', key_hex, '-iv', iv_hex,
            '-in', path_in, '-out', path_out
        ]
        p = subprocess.Popen(cmd_aes, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        out, err = p.communicate()
        if p.returncode != 0:
            raise RuntimeError("OpenSSL AES encrypt error (rc=%d): %s" % (p.returncode, err))

        with open(path_out, 'rb') as f:
            ciphertext = f.read()
    finally:
        try: os.remove(path_in)
        except: pass
        try: os.remove(path_out)
        except: pass

    fd_key_in, path_key_in = tempfile.mkstemp(prefix='rsa_key_in_', dir='/tmp')
    fd_key_out, path_key_out = tempfile.mkstemp(prefix='rsa_key_out_', dir='/tmp')
    os.close(fd_key_in)
    os.close(fd_key_out)

    try:
        with open(path_key_in, 'wb') as f:
            f.write(aes_key)

        cmd_rsa = [
            '/usr/bin/openssl', 'rsautl', '-encrypt',
            '-pubin', '-inkey', pubkey_path,
            '-in', path_key_in, '-out', path_key_out
        ]
        if use_oaep:
            cmd_rsa.insert(3, '-oaep')
        else:
            cmd_rsa.insert(3, '-pkcs')

        p = subprocess.Popen(cmd_rsa, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        out, err = p.communicate()
        if p.returncode != 0:
            raise RuntimeError("OpenSSL rsautl error (rc=%d): %s" % (p.returncode, err))

        with open(path_key_out, 'rb') as f:
            enc_key = f.read()
    finally:
        try: os.remove(path_key_in)
        except: pass
        try: os.remove(path_key_out)
        except: pass

    b64_enc_key = base64.b64encode(enc_key)
    b64_iv = base64.b64encode(iv)
    b64_ct = base64.b64encode(ciphertext)

    return b"HYBRID_V1:" + b64_enc_key + b":" + b64_iv + b":" + b64_ct


def _key_b64_to_hex(key_b64):
    key = base64.b64decode(key_b64)  # 16/24/32 bytes
    try:
        # Py2 cách cũ
        return key.encode('hex')
    except Exception:
        # fallback portable
        return binascii.b2a_hex(key)


def _key_b64_to_hex_and_cipher(key_b64):
    """
    Base64 -> bytes -> (hex_key, cipher_name) theo độ dài key:
    16B -> aes-128-ecb, 24B -> aes-192-ecb, 32B -> aes-256-ecb
    """
    key = base64.b64decode(key_b64)
    klen = len(key)
    if klen == 16:
        cipher = 'aes-128-ecb'
    elif klen == 24:
        cipher = 'aes-192-ecb'
    elif klen == 32:
        cipher = 'aes-256-ecb'
    else:
        raise RuntimeError("Độ dài key không hợp lệ: %d bytes (cần 16/24/32)" % klen)
    # Py2: trả về hex dạng str
    try:
        hexkey = key.encode('hex')
    except Exception:
        hexkey = binascii.b2a_hex(key)
    return hexkey, cipher


def run_bash(script_bytes):
    """Run bash script passed as BYTES via stdin; returns stdout bytes (Py2-safe)."""
    if isinstance(script_bytes, unicode):
        script_bytes = script_bytes.encode('utf-8')
    p = subprocess.Popen("/bin/bash -s", shell=True,
                         stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    out, err = p.communicate(script_bytes)
    if p.returncode != 0:
        # Ghi stderr vào cuối output để tiện debug nhưng không làm hỏng bytes
        out = (out or b"") + b"\n===== STDERR =====\n" + (err or b"")
    return out


ANSI_RE = re.compile(r'\x1b\[[0-9;]*[A-Za-z]')


def strip_ansi(s):
    try:
        # Py2: s có thể là bytes; cứ để bytes, regex vẫn xử được
        return ANSI_RE.sub('', s)
    except Exception:
        return s


def safe_slug(s, maxlen=80):
    # Chuyển sang ASCII-safe cho tên file: thay khoảng trắng/ký tự lạ bằng _
    if isinstance(s, unicode):
        s = s.encode('utf-8')
    # bỏ ANSI, rồi thay ký tự không an toàn
    s = strip_ansi(s)
    s = re.sub(r'[^\w\.-]+', '_', s)  # chỉ giữ a-zA-Z0-9_ . -
    s = s.strip('_')
    if len(s) > maxlen:
        s = s[:maxlen]
    return s


def extract_host_time(text):
    """
    Tìm Hostname và Time bằng regex, không phụ thuộc thứ tự dòng.
    Kỳ vọng trong output có dạng:
      Hostname: myhost
      Time: 2025-11-06 13:00:00
    """
    t = strip_ansi(text)
    # Tách dòng để tìm
    lines = t.splitlines()
    hostname_out, time_out = None, None

    # Regex linh hoạt (bỏ qua hoa/thường, khoảng trắng)
    re_host = re.compile(r'^\s*Hostname:\s*(\S+)\s*$', re.I)
    re_time = re.compile(r'^\s*Time:\s*(.+?)\s*$', re.I)

    for line in lines:
        if hostname_out is None:
            m = re_host.match(line)
            if m:
                hostname_out = m.group(1)
        if time_out is None:
            m = re_time.match(line)
            if m:
                time_out = m.group(1)
        if hostname_out and time_out:
            break

    if not hostname_out:
        # fallback an toàn
        try:
            hostname_out = socket.gethostname()
        except Exception:
            hostname_out = 'unknown-host'

    if not time_out:
        # fallback: now
        try:
            time_out = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
        except Exception:
            time_out = 'unknown-time'

    # Chuẩn hoá
    hostname_out = safe_slug(hostname_out, maxlen=80)
    # format lại time sang safe filename
    time_safe = time_out.replace('/', '-').replace('\\', '-').replace(':', '_').replace(' ', '-')
    time_safe = safe_slug(time_safe, maxlen=80)

    return hostname_out, time_safe


def run_audit(file_path=None):
    # Banner
    print(Fore.BLUE + """
              _    _ _____ _____ _______            _    _          _____  _____  ______ _   _ _____ _   _  _____ 
         /\\  | |  | |  __ \\_   _|__   __|          | |  | |   /\\   |  __ \\|  __ \\|  ____| \\ | |_   _| \\ | |/ ____|
        /  \\ | |  | | |  | || |    | |     ______  | |__| |  /  \\  | |__) | |  | | |__  |  \\| | | | |  \\| | |  __ 
       / /\\ \\| |  | | |  | || |    | |    |______| |  __  | / /\\ \\ |  _  /| |  | |  __| | . ` | | | | . ` | | |_ |
      / ____ \\ |__| | |__| || |_   | |             | |  | |/ ____ \\| | \\ \\| |__| | |____| |\\  |_| |_| |\\  | |__| |
     /_/    \\_\\____/|_____/_____|  |_|             |_|  |_/_/    \\_\\_|  \\_\\_____/|______|_| \\_|_____|_| \\_|\\_____|

    """ + Fore.RESET)
    print(Fore.RED + "Running Audit Hardening..." + Fore.RESET)

    script_path = resolve_script_path(file_path)
    if script_path is None:
        _p(u"ERROR: No audit script found. Use -p <script.sh> or bundle the script.", getattr(Fore, 'RED', None))
        return

    _p(u"  Script: %s" % script_path)
    print("  [1/3] Reading script...")
    with open(script_path, 'r') as f:
        script_content = f.read()

    content_u = script_content
    script_parts = content_u.split(
        "##################################################################################################################")
    output = ""
    for part in script_parts:
        # Chuẩn hoá: part_bytes là bytes UTF-8; so sánh shebang ở dạng bytes
        try:
            unicode
            if isinstance(part, unicode):
                part_bytes = part.encode('utf-8')
            else:
                part_bytes = part
        except NameError:
            part_bytes = part  # (không dùng trên Py3 ở đây)

        if b"#!/bin/bash" not in part_bytes:
            bash_src = b"#!/bin/bash\n" + part_bytes
        else:
            bash_src = part_bytes

        out = run_bash(bash_src)  # out là bytes
        try:
            output += out.decode('utf-8', 'ignore')
        except Exception:
            output += out  # worst-case, giữ nguyên bytes

    # Extract hostname/time lines from the produced output (keep logic as original)
    hostname_out, time_generate = extract_host_time(output)

    os_name_raw = detect_os_name()
    os_name = safe_slug(os_name_raw, maxlen=60)  # ASCII-safe
    file_encrypt_name = '%s_%s_%s.txt.enc' % (hostname_out, os_name, time_generate)
    enc_b64 = write_content(output)

    with open(file_encrypt_name, 'wb') as f:
        # enc_b64 is bytes in Py2
        f.write(enc_b64)

    if is_file_exist(file_encrypt_name):
        _p(u"THÀNH CÔNG - FILE ENCRYPT %s" % file_encrypt_name, getattr(Fore, 'GREEN', None))
    else:
        _p(u"THẤT BẠI!!!", getattr(Fore, 'RED', None))


def run_ubuntu_audit(file_path=None):
    old_file_path = '%s.txt' % hostname
    if is_file_exist(old_file_path):
        try:
            os.remove(old_file_path)
        except Exception:
            pass
    run_audit(file_path)


def parse_options():
    if not OptionParser:
        raise SystemExit("[ERROR] optparse not available. Please use Python 2.6/2.7.")
    parser = OptionParser(usage="usage: %prog [-p PATH_TO_SCRIPT]")
    parser.add_option("-p", "--path", dest="path",
                      help="Path to audit script (optional if bundled)", metavar="FILE", action="store", default=None)
    (options, args) = parser.parse_args()
    return options


if __name__ == '__main__':
    import sys

    opts = parse_options()
    run_ubuntu_audit(opts.path)