From 2e41f34a3e343e3c547c1e228de832a343b95e0b Mon Sep 17 00:00:00 2001
From: Max R <maxr@outlook.com>
Date: Tue, 2 Mar 2021 10:54:00 -0500
Subject: [PATCH] Identify sockets

---
 README.md              |  2 +-
 identify/identify.py   | 16 ++++++++++++----
 tests/identify_test.py | 14 ++++++++++++++
 3 files changed, 27 insertions(+), 5 deletions(-)

diff --git a/README.md b/README.md
index accf29b..68a6893 100644
--- a/README.md
+++ b/README.md
@@ -37,7 +37,7 @@ If you have an actual file on disk, you can get the most information possible
 
 When using a file on disk, the checks performed are:
 
-* File type (file, symlink, directory)
+* File type (file, symlink, directory, socket)
 * Mode (is it executable?)
 * File name (mostly based on extension)
 * If executable, the shebang is read and the interpreter interpreted
diff --git a/identify/identify.py b/identify/identify.py
index 0f936a8..51c1288 100644
--- a/identify/identify.py
+++ b/identify/identify.py
@@ -1,6 +1,7 @@
 import os.path
 import re
 import shlex
+import stat
 import string
 import sys
 from typing import IO
@@ -18,13 +19,14 @@
 
 DIRECTORY = 'directory'
 SYMLINK = 'symlink'
+SOCKET = 'socket'
 FILE = 'file'
 EXECUTABLE = 'executable'
 NON_EXECUTABLE = 'non-executable'
 TEXT = 'text'
 BINARY = 'binary'
 
-TYPE_TAGS = frozenset((DIRECTORY, FILE, SYMLINK))
+TYPE_TAGS = frozenset((DIRECTORY, FILE, SYMLINK, SOCKET))
 MODE_TAGS = frozenset((EXECUTABLE, NON_EXECUTABLE))
 ENCODING_TAGS = frozenset((BINARY, TEXT))
 _ALL_TAGS = {*TYPE_TAGS, *MODE_TAGS, *ENCODING_TAGS}
@@ -36,12 +38,18 @@
 
 
 def tags_from_path(path: str) -> Set[str]:
-    if not os.path.lexists(path):
+    try:
+        sr = os.lstat(path)
+    except (OSError, ValueError):  # same error-handling as `os.lexists()`
         raise ValueError(f'{path} does not exist.')
-    if os.path.isdir(path):
+
+    mode = sr.st_mode
+    if stat.S_ISDIR(mode):
         return {DIRECTORY}
-    if os.path.islink(path):
+    if stat.S_ISLNK(mode):
         return {SYMLINK}
+    if stat.S_ISSOCK(mode):
+        return {SOCKET}
 
     tags = {FILE}
 
diff --git a/tests/identify_test.py b/tests/identify_test.py
index 1fc532e..8cc5856 100644
--- a/tests/identify_test.py
+++ b/tests/identify_test.py
@@ -1,6 +1,8 @@
 import io
 import os
+import socket
 import stat
+from tempfile import TemporaryDirectory
 
 import pytest
 
@@ -12,6 +14,7 @@ def test_all_tags_includes_basic_ones():
     assert 'directory' in identify.ALL_TAGS
     assert 'executable' in identify.ALL_TAGS
     assert 'text' in identify.ALL_TAGS
+    assert 'socket' in identify.ALL_TAGS
 
 
 @pytest.mark.parametrize(
@@ -51,6 +54,17 @@ def test_tags_from_path_symlink(tmpdir):
     assert identify.tags_from_path(x.strpath) == {'symlink'}
 
 
+def test_tags_from_path_socket():
+    tmproot = '/tmp'  # short path avoids `OSError: AF_UNIX path too long`
+    with TemporaryDirectory(dir=tmproot) as tmpdir:
+        socket_path = os.path.join(tmpdir, 'socket')
+        with socket.socket(socket.AF_UNIX) as sock:
+            sock.bind(socket_path)
+            tags = identify.tags_from_path(socket_path)
+
+    assert tags == {'socket'}
+
+
 def test_tags_from_path_broken_symlink(tmpdir):
     x = tmpdir.join('foo')
     x.mksymlinkto(tmpdir.join('lol'))