From e2ab020a8ac18fb217d6152e0e018e48b4365721 Mon Sep 17 00:00:00 2001
From: Anexen <avolk93@gmail.com>
Date: Fri, 26 Jun 2020 14:00:40 +0300
Subject: [PATCH] include package hash in index

---
 pypiprivate/publish.py | 38 +++++++++++++++++++++++++++++++++++---
 pypiprivate/storage.py | 25 +++++++++++++++++++++++++
 2 files changed, 60 insertions(+), 3 deletions(-)

diff --git a/pypiprivate/publish.py b/pypiprivate/publish.py
index 52f8128..42b3270 100644
--- a/pypiprivate/publish.py
+++ b/pypiprivate/publish.py
@@ -1,6 +1,7 @@
 import os
 import re
 import logging
+import hashlib
 
 from pkg_resources import packaging
 from jinja2 import Environment
@@ -21,6 +22,18 @@ def normalized_name(name):
     return re.sub(r"[-_.]+", "-", name).lower()
 
 
+def sha256sum(filename):
+    h = hashlib.sha256()
+    b = bytearray(128 * 1024)
+    mv = memoryview(b)
+
+    with open(filename, 'rb', buffering=0) as f:
+        for n in iter(lambda: f.readinto(mv), 0):
+            h.update(mv[:n])
+
+    return h.hexdigest()
+
+
 def _filter_pkg_dists(dists, pkg_name, pkg_ver):
     # Wheels have different naming conventions: https://www.python.org/dev/peps/pep-0491/#escaping-and-unicode
     # We want to account for both sdist and wheel naming.
@@ -56,7 +69,10 @@ def build_index(title, items, index_type='root'):
     <h1>{{title}}</h1>
     {% endif -%}
     {% for item in items %}
-    <a href="{{item}}">{{item}}</a><br>
+    <a href="{{item["name"]}}#{%- if item.get("sha256") %}sha256={{item["sha256"]}}{% endif -%}">
+        {{item["name"]}}
+    </a>
+    <br>
     {% endfor %}
 </body>
 </html>
@@ -77,11 +93,24 @@ def upload_dist(storage, dist):
     logger.info('Uploading dist: {0}'.format(dist['artifact']))
     dest = storage.join_path(dist['normalized_name'], dist['artifact'])
     storage.put_file(dist['path'], dest, sync=True)
+    content_hash = sha256sum(dist['path'])
+    logger.info('{0} sha256: {1}'.format(dist['artifact'], content_hash))
+    storage.put_contents(content_hash, dest + '.sha256', sync=True)
 
 
 def update_pkg_index(storage, pkg_name):
     logger.info('Updating index for package: {0}'.format(pkg_name))
-    dists = [d for d in storage.listdir(pkg_name) if d != 'index.html']
+    dists = [
+        {
+            "name": d,
+            "sha256": storage.read_contents(
+                storage.join_path(pkg_name, d) + '.sha256',
+                raise_if_not_exist=False
+            )
+        }
+        for d in storage.listdir(pkg_name)
+        if d != 'index.html' and not d.endswith('.sha256')
+    ]
     title = 'Links for {0}'.format(pkg_name)
     index = build_index(title, dists, 'pkg')
     index_path = storage.join_path(pkg_name, 'index.html')
@@ -90,7 +119,10 @@ def update_pkg_index(storage, pkg_name):
 
 def update_root_index(storage):
     logger.info('Updating repository index')
-    pkgs = sorted([p for p in storage.listdir('.') if p != 'index.html'])
+    pkgs = sorted(
+        [{"name": p} for p in storage.listdir('.') if p != 'index.html'],
+        key=lambda x: x["name"]
+    )
     title = 'Private Index'
     index = build_index(title, pkgs, 'root')
     index_path = storage.join_path('index.html')
diff --git a/pypiprivate/storage.py b/pypiprivate/storage.py
index ff5fe64..496b5ec 100644
--- a/pypiprivate/storage.py
+++ b/pypiprivate/storage.py
@@ -36,6 +36,9 @@ def put_contents(self, contents, dest, sync=False):
     def put_file(self, src, dest, sync=False):
         raise NotImplementedError
 
+    def read_contents(self, dest, raise_if_not_exist=True):
+        raise NotImplementedError
+
 
 class LocalFileSystemStorage(Storage):
 
@@ -81,6 +84,16 @@ def put_file(self, src, dest, sync=False):
         shutil.copy(src, dest_path)
         return dest_path
 
+    def read_contents(self, dest, raise_if_not_exist=True):
+        dest_path = self.join_path(self.base_path, dest)
+        try:
+            with open(dest_path, 'r') as f:
+                return f.read()
+        except FileNotFoundError:
+            if raise_if_not_exist:
+                raise
+            return None
+
     def __repr__(self):
         return (
             '<LocalFileSystemStorage(base_path="{0}")>'
@@ -198,6 +211,18 @@ def put_file(self, src, dest, sync=False):
             waiter = client.get_waiter('object_exists')
             waiter.wait(Bucket=self.bucket.name, Key=dest_path)
 
+    def read_contents(self, dest, raise_if_not_exist=True):
+        dest = self.prefixed_path(dest)
+        try:
+            s3_object = self.s3.meta.client.get_object(
+                Bucket=self.bucket.name, Key=dest)
+            return s3_object['Body'].read().decode('utf-8')
+        except ClientError:
+            if raise_if_not_exist:
+                raise FileNotFoundError(
+                    'No such file or directory: {0}'.format(dest))
+            return None
+
     def __repr__(self):
         return (
             '<AWSS3Storage(bucket="{0}", prefix="{1}")>'