1 files changed, 371 insertions, 0 deletions
diff --git a/src/paste.py b/src/paste.py
new file mode 100755
index 0000000..baef4fc
--- /dev/null
+++ b/src/paste.py
@@ -0,0 +1,371 @@
+#!/usr/bin/env python3
+
+import os
+import sys
+import io
+import secrets
+import sqlite3
+import zstandard
+from datetime import datetime
+from urllib.parse import parse_qs
+from pathlib import Path
+from dataclasses import dataclass
+from typing import Optional, Dict, List, Any
+from boto3.session import Session
+from botocore.client import Config
+from botocore.handlers import set_list_objects_encoding_type_url
+from pygments import highlight
+from pygments.lexers import guess_lexer, get_lexer_by_name
+from pygments.formatters import HtmlFormatter
+
+
+@dataclass
+class Environment:
+    request_method: str
+    content_length: int
+    path: Optional[str]
+    parameters: Dict[str, List[str]]
+    content_length_max_bytes: int
+    access_key: str
+    secret_key: str
+    endpoint: str
+    region: str
+    bucket: str
+    signature_version: str
+    key_length: int
+    site: str
+    dict: Path
+    db_path: Path
+
+    def __init__(self, environment: Dict[str, str]):
+        self.request_method = environment["REQUEST_METHOD"]
+        self.content_length = int(environment["CONTENT_LENGTH"])
+
+        try:
+            self.path = environment["PATH_INFO"]
+        except KeyError:
+            self.path = None
+
+        try:
+            query_string = environment["QUERY_STRING"]
+            self.parameters = parse_qs(query_string)
+        except KeyError:
+            self.parameters = {}
+
+        self.content_length_max_bytes = int(
+            environment["PASTE_CONTENT_LENGTH_MAX_BYTES"]
+        )
+        self.access_key = environment["PASTE_ACCESS_KEY"]
+        self.secret_key = environment["PASTE_SECRET_KEY"]
+        self.endpoint = environment["PASTE_ENDPOINT"]
+        self.region = environment["PASTE_REGION"]
+        self.bucket = environment["PASTE_BUCKET_NAME"]
+        self.signature_version = environment["PASTE_SIGNATURE_VERSION"]
+        self.db_path = Path(environment["PASTE_DB"])
+        self.key_length = int(environment["PASTE_KEY_LENGTH"])
+        self.site = environment["PASTE_SITE"]
+        self.dict = Path(environment["PASTE_DICT"])
+        self.storage_max_bytes = int(environment["PASTE_STORAGE_MAX_BYTES"])
+
+
+def connect_to_s3(
+    endpoint: str,
+    region: str,
+    signature_version: str,
+    access_key: str,
+    secret_key: str,
+):
+    session = Session(
+        aws_access_key_id=access_key,
+        aws_secret_access_key=secret_key,
+        region_name=region,
+    )
+
+    session.events.unregister(
+        "before-parameter-build.s3.ListObjects", set_list_objects_encoding_type_url
+    )
+
+    return session.resource(
+        "s3",
+        endpoint_url=endpoint,
+        config=Config(signature_version=signature_version),
+    )
+
+
+def key_in_db(cursor: sqlite3.Cursor, key: str) -> bool:
+    results = cursor.execute("select * from pastes where pastes.key=?", (key,))
+    return results.fetchone() is not None
+
+
+def get_syntax_from_db(cursor: sqlite3.Cursor, key: str) -> Optional[str]:
+    results = cursor.execute(
+        "select pastes.syntax from pastes where pastes.key=?", (key,)
+    )
+    return results.fetchone()
+
+
+def get_storage_use(cursor: sqlite3.Cursor) -> Optional[int]:
+    results = cursor.execute("select sum(pastes.size) from pastes")
+
+    if (row := results.fetchone()) is not None:
+        return row[0]
+    else:
+        return None
+
+
+def delete_oldest_paste(cursor: sqlite3.Cursor) -> Optional[str]:
+    results = cursor.execute(
+        "select pastes.key from pastes order by pastes.datetime limit 1"
+    )
+
+    if (row := results.fetchone()) is not None:
+        key = row[0]
+
+        cursor.execute("delete from pastes where pastes.key=?", (key,))
+
+        return key
+    else:
+        return None
+
+
+def generate_key(words: List[str], length: int) -> str:
+    choices = []
+    for _ in range(length):
+        choices.append(secrets.choice(words))
+
+    return "-".join(word for word in choices).lower()
+
+
+def pygmentize(content: str, syntax: Optional[str]) -> str:
+    if syntax is not None:
+        try:
+            lexer = get_lexer_by_name(syntax)
+        except Exception as e:
+            print(e, file=sys.stderr)
+            lexer = guess_lexer(content)
+    else:
+        lexer = guess_lexer(content)
+
+    formatter = HtmlFormatter(full=True, cssclass="source")
+
+    return highlight(content, lexer, formatter)
+
+
+def main():
+    try:
+        environment = Environment({key: val for key, val in os.environ.items()})
+    except Exception as e:
+        print("Status: 500\r\n")
+        print(f"failed to load environment: {e}", file=sys.stderr)
+        return
+
+    if environment.request_method not in ["GET", "POST"]:
+        print("Status: 400\r\n")
+        print(f"unsupported request method: {environment.request_method}")
+        return
+
+    try:
+        words = environment.dict.read_text().split("\n")
+    except Exception as e:
+        print("Status: 500\r\n")
+        print(f"failed to load dictonary: {e}", file=sys.stderr)
+        return
+
+    try:
+        s3 = connect_to_s3(
+            environment.endpoint,
+            environment.region,
+            environment.signature_version,
+            environment.access_key,
+            environment.secret_key,
+        )
+    except Exception as e:
+        print("Status: 500\r\n")
+        print(f"failed to connect to s3: {e}", file=sys.stderr)
+        return
+
+    try:
+        bucket = s3.Bucket(environment.bucket)
+    except Exception as e:
+        print("Status: 500\r\n")
+        print(f"failed to connect to bucket: {e}", file=sys.stderr)
+        return
+
+    try:
+        connection = sqlite3.connect(environment.db_path)
+        cursor = connection.cursor()
+    except Exception as e:
+        print("Status: 500\r\n")
+        print(f"failed to connect to database: {e}")
+        return
+
+    match environment.request_method:
+        case "POST" if (
+            environment.content_length > environment.content_length_max_bytes
+        ):
+            print("Status: 400\r\n")
+            print(
+                f"content length exceeded maximum length of {environment.content_length_max_bytes}"
+            )
+            return
+
+        case "POST":
+            # delete pastes to make size for new one
+            while (
+                (storage_use := get_storage_use(cursor)) is not None
+                and storage_use + environment.content_length
+                > environment.storage_max_bytes
+            ):
+                match delete_oldest_paste(cursor):
+                    case str(key):
+                        try:
+                            s3.Object(environment.bucket, key).delete()
+                        except Exception as e:
+                            print("Status: 500\r\n")
+                            print(
+                                f"failed to delete object: {key}: {e}", file=sys.stderr
+                            )
+                            return
+
+                        try:
+                            connection.commit()
+                            print(
+                                f"rotated {key} out to make room for new paste",
+                                file=sys.stderr,
+                            )
+                        except Exception as e:
+                            print("Status: 500\r\n")
+                            print(f"failed to commit to database: {e}", file=sys.stderr)
+                            return
+                    case None:
+                        break
+
+            try:
+                syntax = environment.parameters["syntax"][0]
+            except KeyError:
+                syntax = None
+            except IndexError:
+                syntax = None
+            except Exception as e:
+                print("Status: 500\r\n")
+                print(f'failed to load query parameter "syntax": {e}', file=sys.stderr)
+                return
+
+            try:
+                key = generate_key(words, environment.key_length)
+                while key_in_db(cursor, key):
+                    key = generate_key(words, environment.key_length)
+            except Exception as e:
+                print("Status: 500\r\n")
+                print(f"failed to query database: {e}", file=sys.stderr)
+                return
+
+            content = sys.stdin.buffer.read(environment.content_length)
+
+            try:
+                content.decode()
+            except Exception as e:
+                print("Status: 400\r\n")
+                print(f"failed to decode content: {e}")
+                return
+
+            compressed = zstandard.compress(content)
+            buffer = io.BytesIO(compressed)
+
+            try:
+                cursor.execute(
+                    "insert into pastes values (?, ?, ?, ?)",
+                    (
+                        key,
+                        len(compressed),
+                        datetime.now().isoformat(),
+                        syntax,
+                    ),
+                )
+            except Exception as e:
+                print("Status: 500\r\n")
+                print(f"failed to insert into database: {e}", file=sys.stderr)
+                return
+
+            try:
+                bucket.upload_fileobj(buffer, key)
+            except Exception as e:
+                print("Status: 500\r\n")
+                print(f"failed to upload paste: {e}", file=sys.stderr)
+                return
+
+            try:
+                connection.commit()
+            except Exception as e:
+                print("Status: 500\r\n")
+                print(f"failed to commit to db: {e}")
+                return
+
+            print("Status: 200\r\n")
+            print(f"{environment.site}/{key}")
+            return
+
+        case "GET" if environment.path is None:
+            print("Status: 400\r\n")
+            print("please provide a key to fetch")
+            return
+
+        case "GET":
+            # strip leading / from path
+            key = environment.path[1:]
+
+            try:
+                exists = key_in_db(cursor, key)
+            except Exception as e:
+                print("Status: 500\r\n")
+                print(f"failed to query db: {e}")
+                return
+
+            if not exists:
+                print("Status: 404\r\n")
+                return
+
+            buffer = io.BytesIO()
+
+            try:
+                bucket.download_fileobj(key, buffer)
+            except Exception as e:
+                print("Status: 500\r\n")
+                print(f"failed to download from bucket: {e}")
+                return
+
+            buffer.seek(0)
+
+            decompressed = zstandard.decompress(buffer.read())
+            content = decompressed.decode()
+
+            try:
+                syntax = environment.parameters["syntax"][0]
+            except KeyError:
+                try:
+                    syntax = get_syntax_from_db(cursor, key)
+                except Exception as e:
+                    print("Status: 500\r\n")
+                    print(f"failed to query database: {e}")
+                    return
+
+            try:
+                if environment.parameters["raw"]:
+                    raw = True
+                else:
+                    raw = False
+            except KeyError:
+                raw = False
+
+            if raw:
+                print("Status: 200\r\n")
+                print(content)
+                return
+            else:
+                highlighted = pygmentize(content, syntax)
+                print("Status: 200\r\n")
+                print(highlighted)
+                return
+
+
+main()